[autumnust]

Problem identified

Whole-object serialization and deep copies lose the pinned CPU storage of fitted ICL ensemble caches. A loaded model therefore stages cache transfers from pageable memory, losing the efficient transfers and prefetch overlap established by fit. This affects ordinary saved/loaded fitted models through the SDM API.

Current status

The problem is present on current main, 0e148ccc9c816f4c6d885fe88ad38ee3760a86c5, which is also the measured baseline. Freshly fitted CUDA ensembles already pin their caches and prefetch them; the allocation property is lost on serialization. This change benefits newly serialized fitted state. Previously saved objects without the new metadata retain their existing loading behavior, and direct fit/predict without serialization is outside the performance claim.

Solution

Record which fitted batch caches were pinned and restore their CPU pinning during deserialization when CUDA is available. Reuse already-pinned caches, preserve explicit load-device placement, and retain the existing stream reset and cache semantics. The production change is 34 lines in ICLModel; it uses cache state and device placement without benchmark identifiers or workload-specific settings.

Restoration copies existing CPU cache values during loading and may temporarily increase host-memory use. A separate Bioresponse diagnostic measured 0.319 s to restore 13,074,472,032 bytes (12.18 GiB) of cache storage. This loading cost is outside the benchmark's prediction timer and is reported separately. Local checks pass: 71 tests, 2 MPS skips, plus all pre-commit checks. Regression coverage includes pickle, torch save/load, deep/shallow copies, CPU-only loading, explicit CUDA mapping, exact predictions, and clear/refit.

Quantifiable improvement

21.97% lower aggregate prediction latency across all six cases, using the geometric mean of per-case median paired candidate/baseline ratios. The initial four-case suite improves 25.58%. Each case has three complete alternating AB/BA/AB pairs; every pair improves every case. Full prediction arrays are exact and official metric errors are unchanged.

Kumo-Tabular Large, 16 estimators, FP16 autocast, unchanged context/recipe/batching settings; one RTX PRO 6000 Blackwell Server Edition GPU with two fixed worker CPUs. The upstream lifecycle is validation inference, full cached refit, and formal prediction. Table timings are medians; memory values are maxima over the three runs. Reduction percentages use paired ratios; ranges are observed spreads, not confidence intervals.

Case	Prediction ms, baseline → candidate	Paired reduction, median [range]	Fit s, baseline → candidate	Peak GPU MiB, B / C	Peak host GiB, B / C
airfoil_self_noise (initial)	65.99 → 58.66	11.1% [11.1–12.2%]	1.91 → 1.91	3439 / 3829	3.39 / 3.08
MIC (initial)	345.36 → 240.56	30.3% [29.4–31.6%]	8.38 → 8.40	4155 / 4155	16.32 / 16.31
coffee_rating_prediction (initial)	433.26 → 251.23	42.0% [41.4–43.1%]	20.50 → 20.58	7661 / 7661	28.50 / 28.47
Bioresponse (initial)	2409.09 → 2055.74	14.7% [14.6–14.9%]	56.01 → 55.10	10209 / 10209	64.00 / 52.64
credit-g (held out)	81.30 → 67.66	15.6% [14.4–19.6%]	2.49 → 2.51	5587 / 5587	4.89 / 4.30
polish_companies_bankruptcy (held out)	536.44 → 467.56	12.8% [2.7–13.2%]	13.11 → 13.13	5719 / 5719	13.31 / 11.07

Bioresponse saves 353 ms between median prediction times. Polish Companies Bankruptcy has a wider observed reduction range (2.7–13.2%); three pairs do not establish a precise effect size for it. The held-out cases were selected before development of this fix.

The fixed GPU allowance is 81,569,985,331 bytes (75.97 GiB), with a 64 GiB host limit and zero swap. All runs meet the resource policy and verified worker cleanup. GPU values are sampled device usage including driver-reserved memory; framework allocated/reserved peaks are retained separately. One Airfoil candidate run has a 390 MiB higher sampled peak while its framework peaks are unchanged. Baseline Bioresponse reaches the host limit with reclaim but no OOM; the candidate peaks at 52.64 GiB. Restoration still has a temporary host-copy cost; these observations do not imply a universal host-memory reduction.

A separate Nsight trace preserves the same workload, predictions, 1,008 cache copies, and 13.1 GB transferred. Kernel/copy overlap rises from 0 to 190 ms; cache-copy duration falls from 362 to 231 ms, while kernel union time stays near 1.82 s. Profiler timings are excluded from the official latency aggregate.

Reproduction and evidence

SDM candidate: 0e964ee793e9e58239c81e708062f51d243cba72; baseline: 0e148ccc9c816f4c6d885fe88ad38ee3760a86c5. PyTorch 2.14.1+cu130; TabArena a334f8e33d337a3c51df18ec3ab35c37d5e37256; Kumo checkpoint v1.0.1 at 3c3e10bbdb590ace29e7026847f92db3c603096d.

Measurement sources and complete artifacts are retained on 4u8g-tur-0037 under /raid/leisun/repos/structured_dm/execution-notes/inference_efficiency_auto_research/experiments/sdgm-11-cache-pinning/. reproduction.md contains the full guarded command and source-bundle restoration instructions. runs/000008 contains the final initial-suite pairs; held-out comparisons are runs/000009/results/{0,1} plus runs/000010/results/0; diagnostic capture/analysis are runs 11/12. The incomplete third pair in run 9 is excluded after a completion-file race; it remains preserved. The replacement uses a tested atomic-publication fix after worker exit, with model execution and timers unchanged.

From the recorded paired-run harness, with the same qualified environment and resource policy:

python -m measurement.kumo_tabular_arena.paired_runs \
  --baseline "$SDM_BASELINE" --candidate "$SDM_CANDIDATE" \
  --harness "$ARENA_HARNESS" --policy "$ARENA_POLICY" \
  --datasets airfoil_self_noise MIC coffee_rating_prediction Bioresponse \
  --pairs 3 --output "$FRESH_OUTPUT"

Repeat with --datasets credit-g polish_companies_bankruptcy for the held-out cases. Official harness: d554e2676e6860cf3ca012cb486e12fff6a87465, with post-worker publication fix 2b2051b631d0dbeb94a5c13bc89229ea9715dadc for the replacement; orchestration: 40ff0ead30a372073dc904f668a1dbc215cf8b77; diagnostic harness: 00f24138ef9b47071c3d7df7032456f7567aed3b. Verified Git bundles preserve these sources alongside exact commands, identities, prediction arrays, and resource samples. The policy is machine-specific; qualify a new host before reusing the experiment there.

Appendix: standalone Bioresponse reproduction

This is a single-case reproducer for the cache-pinning improvement. It uses the public SDM benchmark adapter and public TabArena runner directly; no private measurement package, research recorder, or Docker wrapper is required. TabArena owns the original validation, full cached refit, fitted-model persistence, warm-up, preprocessing, and fit/prediction timer boundaries. Kumo-Tabular Large uses 16 estimators and FP16 autocast through the pinned adapter's defaults. The selected task is Bioresponse, OpenML task 363620, outer fold 0 / repeat 0, TabArena lite selection.

The comparison is between the exact original measured SDM revisions:

• Baseline: 0e148ccc9c816f4c6d885fe88ad38ee3760a86c5.
• Candidate: 0e964ee793e9e58239c81e708062f51d243cba72.
• TabArena and BenchEval: a334f8e33d337a3c51df18ec3ab35c37d5e37256.
• NVIDIA Kumo-Tabular checkpoint: 3c3e10bbdb590ace29e7026847f92db3c603096d (the originally measured v1.0.1 snapshot).

This does not benchmark the later main-merge commit or the proposed Cache refactor. It does not reproduce the six-case 21.97% aggregate by itself. The original Bioresponse observation was a 14.7% median paired prediction-latency reduction; hardware and software differences can change the magnitude. Fresh fit followed directly by predict is not the affected lifecycle: the fitted cache must undergo serialization/loading.

Setup and execution

Save the three files below together. Use Linux, Python 3.12, one otherwise-idle CUDA GPU, and sufficient host RAM. Validation used an RTX PRO 6000 Blackwell Server Edition, PyTorch 2.14.1+cu130, CUDA 13.0, two fixed CPU cores, a 64 GiB worker memory cap with zero swap, and a fixed 75.97 GiB total GPU allowance. Original Bioresponse runs reached approximately 10 GiB sampled GPU usage and up to 64 GiB worker host memory; leave additional RAM for the operating system. Other GPUs must support the pinned CUDA build and have enough memory; these observations are not a universal minimum-memory guarantee.

python3.12 -m venv .venv
. .venv/bin/activate
python -m pip install -r requirements.txt
# If using uv instead of pip, add --index-strategy unsafe-best-match:
# both indexes in requirements.txt are official, and this matches pip selection.

git clone https://github.com/NVIDIA/structured-data-models.git sdm-source
git -C sdm-source worktree add --detach ../sdm-baseline \
  0e148ccc9c816f4c6d885fe88ad38ee3760a86c5
git -C sdm-source worktree add --detach ../sdm-candidate \
  0e964ee793e9e58239c81e708062f51d243cba72

# Choose an idle physical GPU before running; expose only that GPU.
export CUDA_VISIBLE_DEVICES=0
export XDG_CACHE_HOME="$PWD/cache"
export HF_HOME="$PWD/cache/huggingface"
export TABARENA_CACHE="$PWD/cache/tabarena"
export DATA_FOUNDRY_CACHE="$PWD/cache/data-foundry"

# Download public inputs before timing. No fitted model is shared across runs.
python - <<'PY'
import openml
from huggingface_hub import hf_hub_download

task = openml.tasks.get_task(363620, download_splits=True)
task.get_dataset().get_data(dataset_format="dataframe")
hf_hub_download(
    "nvidia/Kumo-Tabular", "large/classifier.pt",
    revision="3c3e10bbdb590ace29e7026847f92db3c603096d",
)
PY

# Freeze one common allocator allowance before either revision loads a model.
# Reserve 20% of initially free VRAM, plus 2 GiB for non-allocator GPU usage.
GPU_BYTES=$(python - <<'PY'
import torch
free, _ = torch.cuda.mem_get_info()
print(int(0.8 * free) - 2 * 1024**3)
PY
)
CPUS=$(python - <<'PY'
import os
print(",".join(map(str, sorted(os.sched_getaffinity(0))[:2])))
PY
)

taskset -c "$CPUS" python benchmark_cache_pinning.py \
  --baseline "$PWD/sdm-baseline" \
  --candidate "$PWD/sdm-candidate" \
  --dataset Bioresponse --pairs 3 \
  --gpu-bytes "$GPU_BYTES" \
  --output "$PWD/bioresponse-comparison"

Use a fresh output directory. The script verifies the SDM hashes and clean tracked source, launches fresh workers in AB/BA/AB order, and refuses a pair with differing predictions, labels, split hashes, checkpoint hashes, metric errors, or package versions. summary.json reports each pair's fit and prediction seconds, median paired latency reduction, and observed range. Per-worker result.json records the runtime and framework GPU allocation/reservation peaks; logs and prediction arrays remain beside it. The upstream prediction timer excludes serialization/loading and cache repinning, so this is not a claim about total save/load/predict time.

The script's allocator limit and 900-second worker timeout are portable safeguards, not replacements for a process-tree host-memory cap or a device-wide VRAM monitor. Our validation retained those external protections. Run on an idle GPU, provision host memory as above, and use your environment's process-tree limits for constrained machines. OpenML/Hub access is required on first use; downloads can be retried using the same input-cache directory without reusing model results.

Validation of this standalone script

Validated on October 21, 2026 in a fresh environment installed from the files below (137 packages; dependency consistency check passes). Three alternating AB/BA/AB Bioresponse pairs produced a 14.98% median paired prediction-latency reduction, with an observed 14.76–15.01% range. Median prediction times were 2.423 s baseline / 2.060 s candidate; median fit times were 56.18 s / 55.77 s. All six prediction arrays exactly match the retained original Bioresponse predictions, and metric error, checkpoint hashes, and test-row identities match. These are new single-case reproduction measurements, separate from the original six-case aggregate.

Validation peaked at 10,209 MiB sampled GPU usage and reached the 64 GiB process-tree host-memory cap with reclaim, with zero swap, no OOM kills, and verified worker cleanup. It used pre-cached public OpenML inputs because an initial external download stalled before model execution; a separate initial launch failed because the script was not mounted. Neither failed startup contributes timing results. The script itself is independent of the local validation wrapper.

benchmark_cache_pinning.py

"""Reproduce PR #1152 on one TabArena dataset using public upstream adapters."""

import argparse
import gzip
import hashlib
import importlib.metadata
import json
import os
import pickle
import statistics
import subprocess
import sys
from pathlib import Path

BASELINE = "0e148ccc9c816f4c6d885fe88ad38ee3760a86c5"
CANDIDATE = "0e964ee793e9e58239c81e708062f51d243cba72"
CHECKPOINT = "3c3e10bbdb590ace29e7026847f92db3c603096d"


def write_json(path, value):
    path.write_text(json.dumps(value, indent=2, default=str) + "\n")


def git(path, *args):
    return subprocess.check_output(
        ["git", "-c", f"safe.directory={path}", "-C", str(path), *args],
        text=True,
    ).strip()


def worker(args):
    sys.path.insert(0, str(args.sdm))
    import numpy as np
    import pandas as pd
    import torch
    import sdm
    import sdm.models.kumo.tabular.model as model_module
    from benchmark.tabular.model import MODEL_CONFIGS
    from tabarena.benchmark.experiment import TabArenaV0pt1ExperimentBundle
    from tabarena.benchmark.task.wrapper import TaskWrapper
    from tabarena.contexts import TabArenaContext
    from tabarena.utils.config_utils import ConfigGenerator

    assert Path(sdm.__file__).resolve().is_relative_to(args.sdm)
    assert torch.cuda.device_count() == 1, "Expose exactly one GPU"
    torch.set_num_threads(2)
    torch.set_num_interop_threads(2)
    torch.cuda.set_per_process_memory_fraction(
        args.gpu_bytes / torch.cuda.get_device_properties(0).total_memory
    )
    torch.cuda.reset_peak_memory_stats()
    args.output.mkdir(parents=True)
    split_hashes, checkpoint_hashes = [], []
    original_split = TaskWrapper.get_train_test_split
    original_download = model_module.download_checkpoint

    def record_split(self, *positional, **keywords):
        frames = original_split(self, *positional, **keywords)
        split_hashes.append([
            hashlib.sha256(
                pd.util.hash_pandas_object(frame, index=True).values.tobytes()
            ).hexdigest()
            for frame in frames
        ])
        return frames

    def pinned_checkpoint(*positional, **keywords):
        keywords["revision"] = CHECKPOINT
        path = original_download(*positional, **keywords)
        with open(path, "rb") as stream:
            checkpoint_hashes.append(hashlib.file_digest(stream, "sha256").hexdigest())
        return path

    # These hooks record inputs and pin checkpoint identity, not model execution.
    TaskWrapper.get_train_test_split = record_split
    model_module.download_checkpoint = pinned_checkpoint
    config = {
        "kv_cache": False,
        "ag_args_ensemble": {
            "refit_folds": True,
            "fold_fitting_strategy": "sequential_local",
        },
        "ag.refit_hyperparameters": {"kv_cache": True},
    }
    generator = ConfigGenerator(
        search_space={},
        model_cls=MODEL_CONFIGS["kumo-tabular-large"].model_cls,
        manual_configs=[config],
    )
    experiments = TabArenaV0pt1ExperimentBundle(
        models=[(generator, 0)], outer_experiments=False,
    ).build_experiments(num_cpus=2, num_gpus=1, time_limit=600)
    context = TabArenaContext(backend="native")
    jobs = context.build_jobs(
        experiments=experiments, dataset_names=[args.dataset], subset=["lite"],
    )
    assert len(jobs) == 1 and (jobs[0].task.fold, jobs[0].task.repeat) == (0, 0)
    try:
        # Upstream owns validation, cached refit, persistence, warm-up and timers.
        context.run_jobs(
            jobs=jobs, expname=args.output / "raw", register=False, debug_mode=True,
        )
    finally:
        TaskWrapper.get_train_test_split = original_split
        model_module.download_checkpoint = original_download
    paths = list((args.output / "raw/data").glob("**/results.pkl"))
    assert len(paths) == 1
    with gzip.open(paths[0], "rb") as stream:
        result = pickle.load(stream)  # Only this process's own upstream output.
    artifacts = result["simulation_artifacts"]
    predictions = artifacts["pred_proba_dict_test"]
    assert len(predictions) == 1 and split_hashes and checkpoint_hashes
    pred = np.asarray(next(iter(predictions.values())))
    assert pred.size and np.isfinite(pred).all()
    np.save(args.output / "predictions.npy", pred, allow_pickle=False)
    np.save(args.output / "labels.npy", np.asarray(artifacts["y_test"]), allow_pickle=False)
    write_json(args.output / "result.json", {
        "sdm": git(args.sdm, "rev-parse", "HEAD"),
        "dataset": args.dataset, "split_hashes": split_hashes,
        "checkpoint_hashes": sorted(set(checkpoint_hashes)),
        "row_ids": np.asarray(artifacts["y_test_idx"]).tolist(),
        "fit_s": float(result["time_train_s"]),
        "predict_s": float(result["time_infer_s"]),
        "metric_error": float(result["metric_error"]),
        "gpu": torch.cuda.get_device_name(), "torch": torch.__version__,
        "python": sys.version, "cuda": torch.version.cuda,
        "gpu_allocator_limit_bytes": args.gpu_bytes,
        "peak_allocated_bytes": torch.cuda.max_memory_allocated(),
        "peak_reserved_bytes": torch.cuda.max_memory_reserved(),
        "packages": {d.metadata["Name"]: d.version for d in importlib.metadata.distributions()},
    })


def compare(args):
    import numpy as np

    for path, sha in [(args.baseline, BASELINE), (args.candidate, CANDIDATE)]:
        assert git(path, "rev-parse", "HEAD") == sha, f"Expected {sha} at {path}"
        assert not git(path, "status", "--porcelain", "--untracked-files=no"), "Dirty checkout"
    args.output.mkdir(parents=True)  # Refuse existing results.
    env = dict(os.environ, OMP_NUM_THREADS="2", OPENBLAS_NUM_THREADS="2",
               MKL_NUM_THREADS="2", PYTHONHASHSEED="0", SDM_ENABLE_CATEGORY_CHECKS="0")
    records, ratios = [], []
    for pair in range(args.pairs):
        order = ["baseline", "candidate"] if pair % 2 == 0 else ["candidate", "baseline"]
        for side in order:
            folder = args.output / str(pair) / side
            folder.parent.mkdir(exist_ok=True)
            with (folder.parent / f"{side}.log").open("w") as log:
                subprocess.run([
                    sys.executable, str(Path(__file__).resolve()), "--sdm", str(getattr(args, side)),
                    "--dataset", args.dataset, "--gpu-bytes", str(args.gpu_bytes),
                    "--output", str(folder),
                ], env=env, stdout=log, stderr=subprocess.STDOUT, check=True, timeout=900)
        folders = [args.output / str(pair) / side for side in ["baseline", "candidate"]]
        b, c = [json.loads((folder / "result.json").read_text()) for folder in folders]
        for name in ["predictions.npy", "labels.npy"]:
            arrays = [np.load(folder / name, allow_pickle=False) for folder in folders]
            assert arrays[0].dtype == arrays[1].dtype and np.array_equal(*arrays), name
        for name in ["split_hashes", "checkpoint_hashes", "row_ids", "packages", "gpu", "torch"]:
            assert b[name] == c[name], name
        assert c["metric_error"] == b["metric_error"], "Metric changed"
        ratio = c["predict_s"] / b["predict_s"]
        ratios.append(ratio)
        records.append({"pair": pair, "baseline_s": b["predict_s"],
                        "candidate_s": c["predict_s"], "ratio": ratio,
                        "baseline_fit_s": b["fit_s"], "candidate_fit_s": c["fit_s"]})
        write_json(args.output / "summary.json", {
            "dataset": args.dataset, "exact_predictions": True,
            "completed_pairs": len(records), "pairs": records,
            "median_reduction_percent": 100 * (1 - statistics.median(ratios)),
            "reduction_range_percent": [100 * (1 - max(ratios)), 100 * (1 - min(ratios))],
        })
        print(json.dumps(records[-1]), flush=True)
    print((args.output / "summary.json").read_text())


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    for name in ["baseline", "candidate", "sdm", "output"]:
        parser.add_argument("--" + name, type=lambda s: Path(s).resolve(), required=name == "output")
    parser.add_argument("--dataset", default="Bioresponse")
    parser.add_argument("--pairs", type=int, default=3)
    parser.add_argument("--gpu-bytes", type=int, required=True,
                        help="Fixed PyTorch allocator limit, identical for both revisions")
    args = parser.parse_args()
    if args.pairs < 1 or args.gpu_bytes <= 0:
        parser.error("pairs and gpu-bytes must be positive")
    if args.sdm:
        worker(args)
    elif args.baseline and args.candidate:
        compare(args)
    else:
        parser.error("Provide --baseline and --candidate (or internal --sdm)")

requirements.txt

--extra-index-url https://download.pytorch.org/whl/cu130
-c constraints.txt
torch
numpy
pandas
pyarrow
huggingface-hub
safetensors
typing-extensions
autogluon.tabular
autogluon.core[ray]
bencheval @ git+https://github.com/autogluon/tabarena.git@a334f8e33d337a3c51df18ec3ab35c37d5e37256#subdirectory=packages/bencheval
tabarena[preprocessing,data-foundry,search-spaces] @ git+https://github.com/autogluon/tabarena.git@a334f8e33d337a3c51df18ec3ab35c37d5e37256#subdirectory=packages/tabarena

constraints.txt

# Versions recorded for the original PR 1152 measurement environment.
adjusttext==1.4.0
aiohappyeyeballs==2.7.1
aiohttp==3.14.5
aiohttp-cors==0.8.1
aiosignal==1.4.0
annotated-doc==0.0.5
annotated-types==0.8.0
anyio==4.15.1
argon2-cffi==25.1.0
argon2-cffi-bindings==26.1.0
attrs==26.1.0
autogluon.common==1.6.4b20261013
autogluon.core==1.6.4b20261013
autogluon.features==1.6.4b20261013
autogluon.tabular==1.6.4b20261013
autorank==1.2.1
backoff==2.2.1
baycomp==1.0.3
boto3==1.43.113
botocore==1.43.113
catboost==1.2.10
certifi==2026.7.22
cffi==2.1.1
charset-normalizer==3.5.2
click==8.5.0
cloudpickle==3.1.2
colorama==0.4.6
colorful==0.6.0a1
configspace==1.2.2
contourpy==1.4.0
cryptography==50.0.2
cuda-bindings==13.4.3
cuda-pathfinder==1.8.3
cuda-toolkit==13.0.3.0
cycler==0.12.1
data-foundry==0.0.6.dev20260924151450
distlib==0.4.3
distro==1.9.0
einops==0.9.0.dev0
faiss-cpu==1.12.0
filelock==4.1.1
fonttools==4.66.1
formulaic==1.2.2
frozenlist==1.8.0
fsspec==2026.9.0
google-api-core==2.42.0
google-auth==2.62.0
googleapis-common-protos==1.75.5
graphviz==0.21
grpcio==1.85.0rc1
h11==0.16.0
hf-xet==1.7.0
httpcore==1.0.9
httpx==0.28.1
huggingface-hub==1.33.0
idna==3.20
iniconfig==2.3.1
interface-meta==2.0.1
interpret-core==0.7.8
jinja2==3.1.6
jmespath==1.1.0
joblib==1.6.0
jsonschema==4.26.0
jsonschema-specifications==2025.9.1
kiwisolver==1.5.1
liac-arff==2.5.0
lightgbm==4.7.0
lightning-utilities==0.15.3
loguru==0.7.3
markdown-it-py==4.2.0
markupsafe==3.0.4
matplotlib==3.10.1
mdurl==0.1.2
minio==7.2.20
more-itertools==11.2.0
mpmath==1.3.0
msgpack==1.2.3
multidict==7.1.0
narwhals==2.27.1
networkx==3.7
numpy==2.5.4
nvidia-cublas==13.1.1.3
nvidia-cuda-cupti==13.0.85
nvidia-cuda-nvrtc==13.0.88
nvidia-cuda-runtime==13.0.96
nvidia-cudnn-cu13==9.24.0.43
nvidia-cufft==12.0.0.61
nvidia-cufile==1.15.1.6
nvidia-curand==10.4.0.35
nvidia-cusolver==12.0.4.66
nvidia-cusparse==12.6.3.3
nvidia-cusparselt-cu13==0.8.1
nvidia-ml-py==13.615.71
nvidia-nccl-cu13==2.30.7
nvidia-nvjitlink==13.4.92
nvidia-nvshmem-cu13==3.4.5
nvidia-nvtx==13.0.85
omegaconf==2.4.1rc1
opencensus==0.11.4
opencensus-context==0.2.dev0
openml==0.15.1
opentelemetry-api==1.45.1
opentelemetry-exporter-prometheus==0.66b1
opentelemetry-proto==1.45.1
opentelemetry-sdk==1.45.1
opentelemetry-semantic-conventions==0.66b1
packaging==26.3
pandas==2.3.3
patsy==1.0.3
pillow==12.3.0
platformdirs==4.13.0
plotly==7.1.0
pluggy==1.6.0
posthog==7.68.0
prometheus-client==0.26.0
propcache==0.5.4
proto-plus==1.29.0
protobuf==7.36.2
psutil==7.2.2
py-spy==0.4.2
pyarrow==24.0.0
pyasn1==0.6.4
pyasn1-modules==0.4.2
pycparser==3.11
pycryptodome==3.24.1
pydantic==2.14.1
pydantic-core==2.50.1
pydantic-settings==2.15.0
pydot==4.0.1
pygments==2.21.0
pyparsing==3.3.3
pytabkit==1.7.3
pytest==9.1.1
python-dateutil==2.9.0.post0
python-discovery==1.6.2
python-dotenv==1.2.4
pytorch-lightning==2.6.6
pytz==2026.5
pyyaml==6.0.3
ray==2.58.0
referencing==0.37.0
regex==2026.9.29
requests==2.34.2
rich==15.0.0
rpds-py==2026.9.1
rtdl-num-embeddings==0.0.12
ruff==0.17.0
s3transfer==0.19.2
safetensors==0.9.0rc1
scikit-learn==1.9.1
scipy==1.18.1
seaborn==0.13.2
sentence-transformers==6.1.0
setuptools==84.0.0
shellingham==1.5.4
six==1.17.0
skrub==0.10.1
smart-open==8.0.3
statsmodels==0.15.0
sympy==1.14.0
tabdpt==1.3.1
tabicl==2.2.0
tabm==0.0.3
tabpfn==9.2.0
tabpfn-common-utils==0.2.23
tabpfn-extensions==0.6.3
tabulate==0.10.0
threadpoolctl==3.7.0
tokenizers==0.23.3
torch==2.14.1+cu130
torchmetrics==1.9.0
tqdm==4.70.1
transformers==5.19.0
triton==3.8.0
tueplots==0.2.5
typer==0.27.3
typing-extensions==4.16.0
typing-inspection==0.4.4
tzdata==2026.5
urllib3==2.8.0
uuid6==2025.0.1
virtualenv==21.14.7
wrapt==2.5.1rc1
xgboost==3.4.1
xmltodict==1.0.4
yarl==1.25.1

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Source mode: Auto
• Repository: NVIDIA/structured-data-models
• Branch: codex/autoresearch-cache-pinning
• File: .coderabbit.yaml
• Commit: 88e8ab0
• Layers: Repository YAML
• Review profile: QUIET
• Plan: Enterprise
• Run ID: a766584f-d7e7-4c66-8961-9e12886e2f47

📥 Commits

Reviewing files that changed from the base of the PR and between 3f2daa9 and 88e8ab0.

📒 Files selected for processing: 1

• sdm/cache.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 9 remain after this review.

📝 Summary

Summary by CodeRabbit

• Bug Fixes
  • Cache serialization now preserves pinned-memory behavior when restoring on systems with CUDA, helping maintain consistent performance after loading saved models.
  • Restored models can continue making predictions and be cleared and refitted as expected, with CPU fallback when CUDA is unavailable.
• Tests
  • Added coverage for model serialization and copying, including prediction consistency and device handling.

Cache serialization now records pinned-tensor positions and restores pinning for eligible CPU tensors when CUDA is available. Tests cover pickle, deepcopy, Torch serialization, device remapping, prediction equivalence, and cleared-model behavior.

Changes: Pinned Cache Serialization

Layer / File(s)	Summary
Record and restore pinned cache state sdm/cache.py	Cache records the traversal positions of pinned tensors during serialization. When CUDA is available, deserialization pins matching CPU tensors, deduplicates replacements, and sets the pinned-memory round threshold to 1 MB before the first replacement.
Validate serialized model behavior test/models/base/test_icl.py	Tests cover pickle, deepcopy, and Torch serialization, including cache pinning, device placement, prediction equivalence, and cleared-model behavior. A CUDA-aware test checks loading on CPU when CUDA is unavailable.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant SerializationAPI
  participant Cache
  participant CUDA
  SerializationAPI->>Cache: serialize and call __getstate__
  Cache-->>SerializationAPI: return state with pinned-tensor positions
  SerializationAPI->>Cache: deserialize and call __setstate__
  Cache->>CUDA: check availability
  Cache->>Cache: pin eligible CPU tensors and replace references

Loading

Merge Risk: ⚪ Minimal · up to 88e8a

This change restores pinned CPU memory for fitted caches after serialization, so loaded models keep the efficient cache transfers. The reviewed concerns about shared storage and PyTorch version compatibility were not supported by the code, and no merge-blocking issue remains.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 16 functions across 5 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly and concisely summarizes the main change: preserving fitted cache pinning during serialization.
Description check		The description is directly related to the changeset and explains the problem, implementation, validation, performance results, and reproduction steps.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches: 💡 1

📝 Generate docstrings: 💡

• Commit to this branch
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[akihironitta]

I opened a quick patch #1226 targeted to this PR. Mind taking a look? Also, any chance we can have a minimal benchmark script in the PR description (raw text or a link to gist) to reproduce the perf gain?

[autumnust]

I opened a quick patch #1226 targeted to this PR. Mind taking a look? Also, any chance we can have a minimal benchmark script in the PR description (raw text or a link to gist) to reproduce the perf gain?

thanks, updated the PR description with a minimal reproducible setup, see "Appendix"

[codecov]

Codecov Report

❌ Patch coverage is 60.00000% with 8 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (main@2d78842). Learn more about missing BASE report.
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
sdm/cache.py	60.00%	8 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1152   +/-   ##
=======================================
  Coverage        ?   82.91%           
=======================================
  Files           ?      173           
  Lines           ?    11191           
  Branches        ?        0           
=======================================
  Hits            ?     9279           
  Misses          ?     1912           
  Partials        ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.