[Callbacks] FEA new hook for beginning of tasks by jeremiedbb · Pull Request #33494 · scikit-learn/scikit-learn

jeremiedbb · 2026-03-09T19:33:08Z

As discussed during the last meeting, this PR proposes an different version of the callback protocol with 4 callbacks

BaseCallback
-setup
-teardown

FitCallback(BaseCallback)
- on_fit_task_end
- on_fit_task_begin

Previoulsy on_fit_begin used to act as a setup and begin of the root task, and on_fit_end used to act as a teardown and end of the root task.

Now all nodes of the context tree call the same hook at the end the task, including the root. In addition, they all call a same hook at the beginning of the task.

Side effect: setup and tear down don't need to be tied to the callback context, so we can make the decorator deal with setup and teardown it self and we don't need to store the context as an attribute of the estimator anymore which is nice.

ping @ogrisel @FrancoisPgm @StefanieSenger @lesteve

sklearn/callback/_callback_context.py

ogrisel

Thanks, here is a first pass of feedback.

sklearn/callback/_base.py

sklearn/callback/_callback_context.py

sklearn/callback/tests/_utils.py

ogrisel · 2026-03-10T10:57:32Z

sklearn/callback/tests/test_callback_context.py

+    assert callback.count_hooks("on_fit_task_begin") == expected_n_tasks
    assert callback.count_hooks("on_fit_task_end") == expected_n_tasks
-    assert callback.count_hooks("on_fit_end") == 1
+    assert callback.count_hooks("fit_teardown") == 1


Those are great assertions but I think it would be nice to also have one about the sequence of events with estimator names and task names.

That would make a stronger test but also be helpful for SLEP reviewers to better understand the dynamics induced by the callback API design choices.

I agree and it's in my todo list :)
but maybe can be done in a separate PR ?

Cross-referencing our previous discussion about this, starting in #33394 (comment).

…upport callbacks

sklearn/callback/tests/_utils.py

StefanieSenger · 2026-03-10T12:47:20Z

sklearn/callback/tests/test_callback_context.py

+    assert callback.count_hooks("on_fit_task_begin") == expected_n_tasks
    assert callback.count_hooks("on_fit_task_end") == expected_n_tasks
-    assert callback.count_hooks("on_fit_end") == 1
+    assert callback.count_hooks("fit_teardown") == 1


Cross-referencing our previous discussion about this, starting in #33394 (comment).

sklearn/callback/tests/test_callback_context.py

jeremiedbb · 2026-03-10T18:07:43Z

Following up a discussion we had with @ogrisel and @FrancoisPgm, I pushed a new commit where the protocol is reworked. Related to this #33494 (comment)

There's now a base protocol for setup and teardown.
Then a protocol for beginning and end of fit tasks.

It keeps the door open for new protocols for beginning and end of function/predict/transform tasks, while allowing a callback to just implement a minimal protocol for its usage (e.g. _BaseCallback + FunctionCallback for a function only callback).

To allow callbacks to do different things during setup/teardown while not having a specific setup and teardown in each protocal, we thought that the setup and teardown hooks would receive the method where the callback context is created. I haven't implemented that yet since we don't have that kind of need yet. We can add this enhancement a bit later.

sklearn/callback/tests/_utils.py

lesteve

I think this looks good, a few comments below.

sklearn/callback/_callback_support.py

sklearn/callback/_progressbar.py

lesteve · 2026-03-11T06:52:00Z

sklearn/callback/_callback_context.py

+            # Only call the `on_fit_task_begin` hook of callbacks that are not
+            # propagated. For propagated callbacks, the hook will be called by the
+            # sub-estimator's root context (both represent the same task).
+            if callback not in getattr(self, "_propagated_callbacks", []):


Out of curiosity, what is the reason to change isinstance check to attribute check? Is there some subtlety involved or it for simplicity?

It is an attr check because we want to detect the callbacks that were propagated from a parent meta-estimator. For auto-propagated callbacks, we still want to call on_fit_task_begin here for the root meta-estimator.

sklearn/callback/tests/test_callback_context.py

sklearn/callback/_base.py

…k-begin-hook

…fit-task-begin-hook

jeremiedbb · 2026-03-12T18:15:54Z

I merged #33522 into this branch to showcase how the new hooks work with the required changes from #33522.

Basically we need to come back to a setup during context initialization in fit and set the context as an attribute to be able to have the teardown done in the finally block of the decorator.

We should therefore merge #33522 before this one

ogrisel

I gave it another pass and it looks good to me as well.

sklearn/callback/tests/test_callback_context.py

FrancoisPgm

A few super minor nitpicks and some docstrings to update but oterwise LGTM thx :)

sklearn/callback/tests/_utils.py

sklearn/callback/_callback_support.py

sklearn/callback/_callback_context.py

…k-begin-hook

jeremiedbb · 2026-03-14T14:52:54Z

Thanks for the reviews. I'm merging this one as well with 2 approvals to be able to update the SLEP.

split setup/teardown and add on_fit_task_begin hook

7678a34

jeremiedbb added No Changelog Needed Callbacks labels Mar 9, 2026

FrancoisPgm reviewed Mar 10, 2026

View reviewed changes

sklearn/callback/_callback_context.py Show resolved Hide resolved

ogrisel reviewed Mar 10, 2026

View reviewed changes

jeremiedbb added 2 commits March 10, 2026 13:31

add estimator to setup/teardown + fix for sub-estimators that don't s…

1475855

…upport callbacks

doc protocol

1282a70

StefanieSenger reviewed Mar 10, 2026

View reviewed changes

jeremiedbb added 2 commits March 10, 2026 14:00

review

7803147

rework protocol

06f0858

github-actions bot added the CI:Linter failure The linter CI is failing on this PR label Mar 10, 2026

jeremiedbb added 2 commits March 10, 2026 18:54

cln

1da9e46

cln

8710ddf

github-actions bot removed the CI:Linter failure The linter CI is failing on this PR label Mar 10, 2026

cln

fa732f3

lesteve reviewed Mar 11, 2026

View reviewed changes

sklearn/callback/tests/_utils.py Outdated Show resolved Hide resolved

lesteve reviewed Mar 11, 2026

View reviewed changes

jeremiedbb added 2 commits March 11, 2026 11:40

Merge remote-tracking branch 'upstream/callbacks' into add-on-fit-tas…

df6ea6a

…k-begin-hook

pass callbacks by reference

fac061e

jeremiedbb mentioned this pull request Mar 11, 2026

[Callbacks] Make clone keep callbacks #33522

Merged

jeremiedbb added 7 commits March 11, 2026 23:43

make progressbar process safe

fc833ab

lint

b2f12f1

lint

0e4d338

reviews

eec9f66

rename uuid

31b5fa4

sanity test for coverage

87b5d97

Merge branch 'there-and-back-again-clone-keeps-callback' into add-on-…

c120393

…fit-task-begin-hook

StefanieSenger mentioned this pull request Mar 13, 2026

[Callbacks] ENH Add callback support to BaseSearchCV #33533

Draft

4 tasks

ogrisel approved these changes Mar 13, 2026

View reviewed changes

sklearn/callback/tests/test_callback_context.py Outdated Show resolved Hide resolved

FrancoisPgm approved these changes Mar 13, 2026

View reviewed changes

jeremiedbb added 3 commits March 13, 2026 17:51

review

96bcc7a

Merge remote-tracking branch 'upstream/callbacks' into add-on-fit-tas…

128bf60

…k-begin-hook

cln

316f44e

jeremiedbb merged commit 66b1eed into scikit-learn:callbacks Mar 14, 2026
25 checks passed

jeremiedbb moved this to Done in Labs Mar 14, 2026

jeremiedbb added this to Labs Mar 14, 2026

jeremiedbb mentioned this pull request Mar 22, 2026

FIX Remove spurious line in progressbar #33599

Merged

Uh oh!

Conversation

jeremiedbb commented Mar 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

ogrisel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ogrisel Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

jeremiedbb Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

StefanieSenger Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

StefanieSenger Mar 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jeremiedbb commented Mar 10, 2026

Uh oh!

Uh oh!

lesteve left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lesteve Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

FrancoisPgm Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jeremiedbb commented Mar 12, 2026

Uh oh!

ogrisel left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

FrancoisPgm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jeremiedbb commented Mar 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

jeremiedbb commented Mar 9, 2026 •

edited

Loading