add caching to .skb.apply

[jeromedockes]

adds caching to estimators added with .skb.apply() and function calls added with deferred() or .skb.apply_func(). Caching is off by default and can be enabled with set_config(cache=True) or set_config(cache='/path/to/cache_dir'). For a given estimator or call, caching can be forbidden entirely in the pipeline definition e.g. X.skb.apply_func(f, no_cache=True) in which case the call is never cached regardless of configuration.

[jeromedockes]

example script

import time

import skrub
from sklearn.datasets import make_classification
from sklearn.ensemble import HistGradientBoostingClassifier

skrub.set_config(cache=True)

X_a, y_a = make_classification(n_samples=50_000, n_features=50)
pred = skrub.X(X_a).skb.apply(HistGradientBoostingClassifier(), y=skrub.y(y_a))
split = pred.skb.train_test_split()
learner = pred.skb.make_learner()

tic = time.perf_counter()
learner.fit(split["train"])
toc = time.perf_counter()
print("first fit", toc - tic)

tic = time.perf_counter()
learner.fit(split["train"])
toc = time.perf_counter()
print("second fit", toc - tic)

tic = time.perf_counter()
result = learner.predict_proba(split["test"])
toc = time.perf_counter()
print("first predict", toc - tic)

tic = time.perf_counter()
result = learner.predict_proba(split["test"])
toc = time.perf_counter()
print("second predict", toc - tic)

first fit 0.28028534399345517
second fit 0.022888949999469332
first predict 0.014127899004961364
second predict 0.005811317008920014

[jeromedockes]

I tried a bit the caching on a realistic pipeline which I took from here:

https://github.com/probabl-ai/skore-example-electricity-load

running the cross-validation with a warm cache we see 10x speedup on the HGB pipeline and 42x on the TabICL one (that is re-running the exact same pipeline though). disclaimer: I didn't check the code that ran the experiments. on that pipeline only the final estimator is costly so caching doesn't help much for hyperparam search for example. with caching disabled the runtime was the same as on main, and on a cold cache the biggest overhead due to writing results was +13% (for the fastest pipeline, HGB + no quantile regression). the outputs were the same as those obtained from the main branch. the cache size did grow quickly, up to 2G, so providing guidance on how to handle that, and later tooling, will be important.

it would be interesting to do some similar experiments on a pipeline that does have expensive intermediate steps, e.g. a tablevectorizer with text

[jeromedockes]

irl we decided we need limiting of the cache size in this PR. I added one possible option but reverting to draft until we have settled on the api and I add the tests

[rcap107]

Do you think this PR will make it in the next release? 🤔

[jeromedockes]

Do you think this PR will make it in the next release? 🤔

i think so 👍

new changes need new review

[rcap107]

Looks good to me, the cache pruning change looks like a fairly light addition overall

[jeromedockes]

cool, thanks for reviewing!

[GaelVaroquaux]

Merged [1]#2017 into main.

Wohoo! Exciting