FIX Don't send ScoringMonitor._log to the pickler (scikit-learn#34821)

Author: jeremiedbb

doc/whats_new/upcoming_changes/callback/34821.fix.rst

@@ -0,0 +1,4 @@
+- Fixed a bug where fitting an estimator with a :class:`callback.ScoringMonitor` and a
+  parallel backend using multi-processing could fail with "Could not pickle the task to
+  send it to the workers".
+  By :user:`Jérémie du Boisberranger <jeremiedbb>`.

sklearn/callback/_scoring_monitor.py

@@ -170,6 +170,12 @@ def on_fit_task_end(
 
         send(self._listener_handle, (run_id, run_info, task_info_path, scores))
 
+    def __getstate__(self):
+        # `_log` is grown by the listener thread, which can run while this callback is
+        # being pickled by another thread, e.g. loky's queue feeder when a task is
+        # dispatched to a worker. The pickler must therefore never walk the live list.
+        return {**self.__dict__, "_log": list(self._log)}
+
     def __setstate__(self, state):
         """Restore state, opening a fresh listener if the inherited one is unusable."""
         self.__dict__.update(state)

sklearn/callback/tests/test_pickle.py

@@ -8,6 +8,8 @@
 unpickled in a fresh Python interpreter.
 """
 
+import functools
+import io
 import pickle
 import re
 import subprocess
@@ -17,6 +19,7 @@
 import pytest
 
 from sklearn.callback import ProgressBar, ScoringMonitor
+from sklearn.callback._transport import _message_consumers
 from sklearn.callback.tests._common.estimators import MaxIterEstimator
 from sklearn.datasets import make_regression
 
@@ -139,3 +142,78 @@ def test_callbacks_refit_after_load_in_fresh_process(tmp_path, capsys):
     stdout = result.stdout.decode()
     assert re.search(r"MaxIterEstimator - fit", stdout)
     assert re.search(r"100%", stdout)
+
+
+class _TraversalRecorder(pickle.Pickler):
+    """A pickler that records the id of every object it walks through.
+
+    `persistent_id` is called for every object the pickler encounters, so an object is
+    recorded whichever path leads to it, not only when it is a direct attribute of the
+    object being pickled,
+    see https://docs.python.org/3/library/pickle.html#pickle.Pickler.persistent_id.
+    """
+
+    def __init__(self):
+        super().__init__(io.BytesIO(), protocol=pickle.HIGHEST_PROTOCOL)
+        self.walked_through = set()
+
+    def persistent_id(self, obj):
+        self.walked_through.add(id(obj))
+        return None  # pickle `obj` as usual
+
+
+def _checked(hook):
+    """Wrap a callback hook so that it first checks the callback it is called on.
+
+    The check is that the pickler does not walk through the objects that listener
+    threads mutate. They are found through the registered consumers: a consumer is
+    normally a method bound to the container it fills, e.g. `self._log.append` for
+    ScoringMonitor or `queue.put` for ProgressBar, so the container is what it is bound
+    to. A consumer bound to nothing, e.g. a closure, is skipped, since there is then no
+    way to tell what it mutates.
+    """
+
+    # preserve the signature of the hook because callbacks are validated against it
+    @functools.wraps(hook)
+    def checked_hook(self, *args, **kwargs):
+        # snapshot because another thread may register a listener concurrently
+        consumers = list(_message_consumers.values())
+        # the containers are held, not just their ids, which could be reused once freed
+        watched = [c.__self__ for c in consumers if hasattr(c, "__self__")]
+
+        recorder = _TraversalRecorder()
+        recorder.dump(self)
+
+        offenders = recorder.walked_through & {id(container) for container in watched}
+        assert not offenders, (
+            f"Pickling {self.__class__.__name__} walks through a container that a"
+            " listener thread mutates concurrently. Keep it away from the pickler,"
+            " either by handing over a copy of it in the callback's `__getstate__`, or"
+            " by storing it outside of the callback instance."
+        )
+
+        return hook(self, *args, **kwargs)
+
+    return checked_hook
+
+
+@pytest.mark.parametrize("factory", CALLBACK_FACTORIES)
+def test_listener_state_is_not_walked_by_the_pickler(factory, monkeypatch):
+    """Check that pickling a callback never traverses state its listener mutates.
+
+    An estimator carrying a callback can be pickled by a background thread, e.g. loky's
+    queue feeder dispatching a task to a worker, while the listener thread of that same
+    callback mutates the callback's state as messages come in. Pickling a container that
+    another thread mutates breaks the dump, which joblib reports as "Could not pickle
+    the task to send it to the workers".
+
+    The check runs from a hook, i.e. while the listeners are up, which is when such a
+    dispatch would happen.
+    """
+    callback = factory()
+
+    for hook_name in ("on_fit_task_begin", "on_fit_task_end"):
+        hook = getattr(callback.__class__, hook_name)
+        monkeypatch.setattr(callback.__class__, hook_name, _checked(hook))
+
+    MaxIterEstimator(max_iter=3).set_callbacks(callback).fit()