Feeding a clanker? Grab this page as raw .md →

FitProgress API#

scikit-learn 1.9 added callbacks: an estimator reports every task it runs, such as each candidate/fold fit of a search and each solver iteration inside it. Its built-in ProgressBar callback draws these with rich, which works well in a terminal but becomes a tall static text block in a notebook. FitProgress draws the same events as a widget.

Display the widget, then fit through it:

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV
from wigglystuff import FitProgress

progress = FitProgress()
progress  # show it in one cell ...

# ... and fit through it in another; you get the fitted search back
grid = progress.fit(
    GridSearchCV(
        LogisticRegression(max_iter=150),
        {"C": [1e-4, 1e-3, 1e-2], "class_weight": [None, "balanced"]},
        cv=5,
    ),
    X,
    y,
)

GridSearchCV and RandomizedSearchCV render as a grid with one row per candidate (one column per hyperparameter) and one column per fold:

A running cell fills up with the inner estimator's progress, for example its lbfgs iterations.
A finished cell shows its test score, shaded darker for better scores. Each row shows its mean once all its folds are in, and when the search is done the best row is highlighted.
A cell outlined in amber used every iteration (max_iter), so the solver may not have converged. A failed fit is outlined in red.
Hover a cell to see its parameters, iterations, fit time, score and rank.
The refit gets one bar at the bottom, with the chosen parameters below it.

Estimators report progress in different ways. A Pipeline reports its steps, and the final estimator inside it can still report iterations; hover a cell to see both. An estimator without callback support, such as a random forest, reports nothing while it fits. Its cells show a striped bar until the fit finishes and then get their score.

fit attaches its callback only for the duration of the fit, so the returned estimator keeps exactly the callbacks it had before. To show test scores while the search runs, it also wraps a single-metric search's scorer during the fit and puts the original scoring and scorer_ back afterwards. A multi-metric search shows its scores (for the refit metric) once it is done. Fits that run in worker processes (n_jobs) report back live too.

Outside a notebook, fit uses scikit-learn's own terminal ProgressBar when rich is installed, and otherwise fits without progress.

Requires scikit-learn 1.9 or later (pip install wigglystuff[sklearn]). The callback API is experimental in scikit-learn, and FitProgress uses its private sklearn.callback._transport module to receive events from worker processes, so a future scikit-learn release may need an update here.

Bases: AnyWidget

Live progress bars for a scikit-learn fit, with a grid view for hyperparameter searches.

Display the widget, then run the fit through it. GridSearchCV and RandomizedSearchCV render as a candidates x folds grid: a running cell fills up with the inner estimator's progress (e.g. solver iterations), a finished cell shows its test score, and each row gets its mean score once all its folds are in. Other estimators render as nested bars.

Outside a notebook, fit uses scikit-learn's own terminal ProgressBar when rich is installed, and fits silently otherwise.

Needs scikit-learn 1.9+ (pip install wigglystuff[sklearn]).

Examples:

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import GridSearchCV
from wigglystuff import FitProgress

progress = FitProgress()
progress  # show it in one cell ...

# ... and fit through it in another
grid = progress.fit(
    GridSearchCV(LogisticRegression(), {"C": [0.01, 0.1, 1.0]}, cv=5),
    X,
    y,
)

Create a FitProgress view.

Parameters:

Name	Type	Description	Default
`max_propagation_depth`	`int \| None`	How many levels of nested estimators report progress. `None` (default) shows every level, e.g. the lbfgs iterations of a LogisticRegression inside a Pipeline inside a search; `1` stops at the search's direct inner estimator; `0` only shows the outer estimator.	`None`
`color`	`str`	Optional CSS color for the bars, the score shading and the best row. Empty uses the default blue.	`''`
`sync_throttle_ms`	`int`	Minimum milliseconds between updates sent to the browser while a fit is running.	`100`
`**kwargs`	`Any`	Forwarded to `anywidget.AnyWidget`.	`{}`

Source code in wigglystuff/fit_progress.py

def __init__(
    self,
    max_propagation_depth: int | None = None,
    color: str = "",
    sync_throttle_ms: int = 100,
    **kwargs: Any,
) -> None:
    """Create a FitProgress view.

    Args:
        max_propagation_depth: How many levels of nested estimators report
            progress. ``None`` (default) shows every level, e.g. the lbfgs
            iterations of a LogisticRegression inside a Pipeline inside a
            search; ``1`` stops at the search's direct inner estimator;
            ``0`` only shows the outer estimator.
        color: Optional CSS color for the bars, the score shading and the
            best row. Empty uses the default blue.
        sync_throttle_ms: Minimum milliseconds between updates sent to the
            browser while a fit is running.
        **kwargs: Forwarded to ``anywidget.AnyWidget``.
    """
    if max_propagation_depth is not None and max_propagation_depth < 0:
        raise ValueError("max_propagation_depth must be >= 0 or None")
    if sync_throttle_ms < 0:
        raise ValueError("sync_throttle_ms must be >= 0")
    super().__init__(color=color, sync_throttle_ms=sync_throttle_ms, **kwargs)
    self.max_propagation_depth = max_propagation_depth
    self._lock = threading.Lock()
    self._nodes: dict[tuple[int, ...], dict[str, Any]] = {}
    self._cell_by_worker: dict[str, tuple[int, ...]] = {}
    self._last_sync = float("-inf")

clear #

clear() -> None

Forget the previous run.

Source code in wigglystuff/fit_progress.py

def clear(self) -> None:
    """Forget the previous run."""
    with self._lock:
        self._nodes = {}
        self._cell_by_worker = {}
    self.tasks = []
    self.candidates = []
    self.scores = None
    self.status = "idle"
    self.started_ms = 0.0
    self.now_ms = 0.0

fit #

fit(estimator: Any, X: Any, y: Any = None, **fit_params: Any) -> Any

Fit estimator while this widget shows its progress; returns the fitted estimator.

The progress callback is only attached for the duration of the fit, so the returned estimator keeps exactly the callbacks it had before.

Source code in wigglystuff/fit_progress.py

def fit(self, estimator: Any, X: Any, y: Any = None, **fit_params: Any) -> Any:
    """Fit ``estimator`` while this widget shows its progress; returns the fitted estimator.

    The progress callback is only attached for the duration of the fit, so
    the returned estimator keeps exactly the callbacks it had before.
    """
    if not hasattr(estimator, "set_callbacks"):
        raise TypeError(
            f"{type(estimator).__name__} does not support scikit-learn callbacks "
            "(needs scikit-learn 1.9+ and an estimator that implements them)"
        )
    if not _in_notebook():
        return self._fit_in_terminal(estimator, X, y, **fit_params)

    from sklearn.callback._transport import close_listener, open_listener

    self.clear()
    param_names = _search_param_names(estimator)
    self.candidates = _grid_candidates(estimator, param_names)
    self.status = "running"
    self.started_ms = time.time() * 1000
    handle = open_listener(self._on_message)
    depth = _propagation_depth(estimator, self.max_propagation_depth)
    callback = _FitProgressCallback(handle, param_names, depth)
    original = list(getattr(estimator, "_skl_callbacks", []))
    estimator.set_callbacks(*original, callback)
    # Single-metric searches report their fold scores live (multi-metric
    # ones only show scores once the search is done).
    scoring = getattr(estimator, "scoring", None)
    reporter = None
    if hasattr(estimator, "cv") and (scoring is None or isinstance(scoring, str) or callable(scoring)):
        from sklearn.metrics import check_scoring

        reporter = _ScoreReporter(check_scoring(estimator.estimator, scoring), handle)
        estimator.scoring = reporter
    try:
        with _no_callback_support_warnings():
            estimator.fit(X, y, **fit_params)
    except BaseException:
        self._finish("error")
        raise
    finally:
        estimator.set_callbacks(*original)
        if reporter is not None:
            estimator.scoring = scoring
            if getattr(estimator, "scorer_", None) is reporter:
                estimator.scorer_ = reporter.scorer
        close_listener(handle)
    if hasattr(estimator, "cv_results_"):
        self.candidates = [_format_params(p, param_names) for p in estimator.cv_results_["params"]]
    self.scores = self._search_scores(estimator)
    self._finish("done")
    return estimator

Synced traitlets#

Traitlet	Type	Notes
`tasks`	`list[dict]`	One dict per task that reports progress: `path` (task ids from the root), `estimator`, `task`, `source_task`, `total`/`done` subtasks, `state`, `start_ms`/`end_ms`, and for search fits `params`, `max_iter`, `unit` (what its subtasks count, e.g. `"lbfgs-iter"`) and the live test `score`.
`candidates`	`list[dict]`	Formatted parameters of each candidate: known up front for `GridSearchCV`, from `cv_results_` once a search is done.
`scores`	`dict \\| None`	After a search: `metric`, `folds` (candidate × fold test scores, `None` for failed fits), `mean`, `rank`, `best_index`.
`color`	`str`	CSS color for the bars, the score shading and the best row; `""` uses the default blue.
`status`	`str`	`"idle"`, `"running"`, `"done"` or `"error"`.
`started_ms`	`float`	Epoch ms when the fit started.
`now_ms`	`float`	Kernel clock at the last update.
`sync_throttle_ms`	`int`	Minimum ms between updates sent to the browser while fitting.