diff --git a/docs/tabicl_strong_feature_comparison.md b/docs/tabicl_strong_feature_comparison.md new file mode 100644 index 0000000..8ad8a4d --- /dev/null +++ b/docs/tabicl_strong_feature_comparison.md @@ -0,0 +1,117 @@ +# GBM (strong) vs TabICL (strong, full capability) on MIMIC + +**Date:** 2026-08-27 + +## 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 report is the first TabICL run on the strong panel at its actual full capability +(`n_estimators=8`, 50,000-row context, no reduction). Getting there took three attempts: + +1. A reduced config (`n_estimators=1`, 20,000-row context) fit and scored fine on the original + 82 GB-RAM host, but it isn't TabICL at its documented best, and Amrit asked for the real thing. +2. Disk offload (`offload_mode="disk"`) was tried next, to fit the full config in less RAM. It + works for fitting, but TabICL rereads its entire context from disk on every `predict_proba` + call with no caching. Scoring 200 rows took 12 minutes; scoring a normal batch (8,192 rows) + did not finish in 15 minutes. Extrapolated, the full comparison would take a day or more. Not + a memory problem, a disk I/O throughput problem. +3. The fix was hardware, not software: the host was migrated to `a2-ultragpu-1g` (170 GB RAM, one + A100-80GB, `odyssey-cbm-a100-ultra`, `us-central1-a`). RAM-resident full capability + (`offload_mode="cpu"`, no disk) fits comfortably (peak RSS ~102 GB) and scores at normal + compute speed, just slower than the reduced config because there is more to compute: roughly + 1,800 seconds (30 minutes) to score one horizon, versus ~120 seconds for the reduced config. + All 12 cells took about 6 hours of sequential GPU compute. + +## Setup + +- Checkpoint: `subset_run_v8_taskset_v3` (MIMIC, task_set v3, `model_kind=cbm`) +- GPU host: `odyssey-cbm-a100-ultra` (`a2-ultragpu-1g`, 170 GB RAM, one A100-80GB) +- 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, full): `n_estimators=8`, 50,000-row context (this project's real defaults, not + reduced), fit on 8 of the 30 train shards, `offload_mode="cpu"`, scored on the same 4 held-out + shards. One model fit, scored, and dropped at a time; no more than one model resident at once. +- Both models are scored on exactly the same rows and labels +- Script: `scripts/tabicl_strong_compare.py` (committed, reusable; reuses the existing GBM scores + instead of refitting) + +**Label caveat, found after this sweep completed, not yet corrected:** the acute kidney injury +label used here, for both the GBM and TabICL scores, comes from `alerts_rows_v3.parquet`, +generated 2026-08-25 05:07 UTC. The KDIGO AKI staging completeness fix (commit `3d7ecbb`, adding +the renal-replacement-therapy and urine-output legs that were previously missing, creatinine-only +staging before that) landed later the same day, at 14:04 UTC. The model checkpoint itself was +also trained before the fix (03:35 UTC). So every AKI number in this report, GBM and TabICL both, +uses the old, incomplete label, which under-counts true AKI-3 cases that only qualify through RRT +or oliguria, not creatinine. This does not explain away the AKI gap below, since the same label +is used for both models, but it means the true AKI numbers (positive or negative for either model) +are not yet known. Worth a rerun with a fresh alerts dump once that's a priority. + +## 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 subsampled context. At full capability, that context is closer to the library's own +validated regime (its authors report strong results up to 50,000 rows), so this asymmetry is +smaller here than it was in the reduced-config run, but it has not gone away. + +## Results + +95% subject-clustered bootstrap confidence intervals, 1000 resamples +(`odyssey.inference.uncertainty.bootstrap_auroc`). + +| Event | Horizon | n | GBM (strong) AUROC | TabICL (strong, full) AUROC | Gap | Verdict | +|---|---|---|---|---|---|---| +| Acute kidney injury | 8h | 95,471 | 0.894 [0.881, 0.906] | 0.766 [0.746, 0.786] | 0.128 | real gap | +| Acute kidney injury | 24h | 84,147 | 0.845 [0.827, 0.861] | 0.739 [0.714, 0.761] | 0.106 | real gap | +| Acute kidney injury | 72h | 57,442 | 0.782 [0.758, 0.805] | 0.685 [0.657, 0.714] | 0.097 | real gap | +| Death | 8h | 136,850 | 0.953 [0.934, 0.969] | 0.961 [0.941, 0.975] | -0.007 | within noise | +| Death | 24h | 135,061 | 0.959 [0.947, 0.969] | 0.940 [0.917, 0.958] | 0.020 | within noise | +| Death | 72h | 130,818 | 0.940 [0.923, 0.954] | 0.925 [0.905, 0.944] | 0.015 | within noise | +| ICU admission | 8h | 85,838 | 0.968 [0.962, 0.974] | 0.963 [0.956, 0.969] | 0.005 | within noise | +| ICU admission | 24h | 74,839 | 0.954 [0.945, 0.963] | 0.940 [0.929, 0.951] | 0.014 | within noise | +| ICU admission | 72h | 48,971 | 0.931 [0.912, 0.946] | 0.914 [0.894, 0.933] | 0.017 | within noise | +| Vasopressor start | 8h | 111,450 | 0.934 [0.916, 0.950] | 0.916 [0.896, 0.936] | 0.018 | within noise | +| Vasopressor start | 24h | 98,722 | 0.914 [0.895, 0.933] | 0.893 [0.871, 0.915] | 0.021 | within noise | +| Vasopressor start | 72h | 67,385 | 0.883 [0.850, 0.913] | 0.868 [0.840, 0.893] | 0.014 | within noise | + +"Real gap" means the two confidence intervals do not overlap. "Within noise" means they do. + +## What this means + +Full capability changes the answer. At the reduced config, TabICL lost on all 12 cells, 7 of them +a real (CI-separated) loss. At full capability, TabICL is statistically indistinguishable from +the tuned GBM on 9 of 12 cells: all of death, all of ICU admission, all of vasopressor start. +Death@8h even slightly favors TabICL, though within noise. + +The reduced config was a real handicap, not a formality. Cutting the ensemble from 8 members to 1 +and the context from 50,000 rows to 20,000 cost TabICL real, measurable performance across the +board, not just on the cells that happened to look weak. + +Acute kidney injury is the one place a real gap survives at full capability (0.097 to 0.128, +separated 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, not from having more context to work +with. TabICL sees the same raw features but has no equivalent way to aggregate them across time +in a single forward pass, at any context size. But see the label caveat above: this comparison +used an AKI label known to be incomplete, so the true size of this gap (bigger, smaller, or +possibly gone) is not yet established. + +## Where things live + +- Comparison script (committed): `scripts/tabicl_strong_compare.py` +- Code changes that made this 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): pulled to + `/tmp/tabicl_full_sweep.json` this session. Ask if you want it moved somewhere durable. +- GBM (strong) scores: `~/runs/subset_run_v8_taskset_v3/alerts_rows_v3.parquet`, generated + 2026-08-25 05:07 UTC (see the label caveat above) +- GPU host used for this run: `odyssey-cbm-a100-ultra` (`a2-ultragpu-1g`, `us-central1-a`), + stopped after this sweep completed. The original host, `odyssey-cbm-a100` (`a2-highgpu-1g`, + `us-central1-f`), is stopped and untouched, kept as a fallback. 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