sklearn/callback/_progressbar.py - main - githubmirror/scikit-learn

sklearn/callback/_progressbar.py
311 строк12 KB
Jérémie du Boisberranger
UI Add milliseconds to elapsed time in progress bars (#34245)
12 июн 2026, 12:54
12 июн 2026, 12:540ae7000
# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause

from datetime import timedelta
from numbers import Integral
from queue import Queue
from threading import Thread

from sklearn.callback._callback_context import get_context_path
from sklearn.callback._transport import close_listener, open_listener, send
from sklearn.utils._optional_dependencies import check_rich_support
from sklearn.utils._param_validation import Interval, validate_params

# Per-fit local queues and monitors, keyed by the run's `root_uuid`. Both are not
# picklable so they live here rather than on the callback instance. Entries are added in
# `setup` and removed in `teardown`.
_run_queues = {}
_run_monitors = {}


class ProgressBar:
    """Callback that displays progress bars for each iterative step of an estimator.

    Parameters
    ----------
    max_propagation_depth : int, default=1
        The maximum depth of nested levels of estimators to display progress bars for.
        0 means that the progress of only the outermost estimator is displayed.
        If set to None, all levels are displayed.

    Notes
    -----
    This callback requires rich to be installed.
    """

    @validate_params(
        {"max_propagation_depth": [Interval(Integral, 0, None, closed="left"), None]},
        prefer_skip_nested_validation=True,
    )
    def __init__(self, max_propagation_depth=1):
        check_rich_support("Progressbar")

        self.max_propagation_depth = max_propagation_depth

        # Handles to the main-process per-fit listeners, keyed by `root_uuid`.
        self._listener_handles = {}

    def setup(self, estimator, context):
        # Lazily create the per-fit transport state. `setup` runs on the main
        # process for the common case where `ProgressBar` is registered on (or
        # auto-propagated from) the outermost estimator. In the degraded case
        # where the outer function does not support callbacks, this may run
        # inside a worker process; that worker will then open its own local
        # listener and the events will be process-local — the callback runs in
        # slightly degraded mode but does not crash.
        queue = Queue()
        # `queue.put` is the message consumer that `send` calls will use to forward
        # information to the rich monitor thread.
        self._listener_handles[context.root_uuid] = open_listener(queue.put)

        progress_monitor = RichProgressMonitor(queue=queue)
        progress_monitor.start()

        _run_queues[context.root_uuid] = queue
        _run_monitors[context.root_uuid] = progress_monitor

    def on_fit_task_begin(self, estimator, context):
        # A new progress bar is created at the beginning of each task that is not a
        # leaf, except if it's also the root task of an estimator.
        if (
            context.max_subtasks != 0
            or context.parent is None
            or context.source_estimator_name is not None
        ):
            # We pass the minimal information to the queue that is necessary to create a
            # progress bar and not the context to avoid pickling the whole context tree.
            path = [ctx.task_id for ctx in get_context_path(context)]
            send(
                self._listener_handles[context.root_uuid],
                {
                    "event": "begin",
                    "path": path,
                    "task_name": context.task_name,
                    "task_id": context.task_id,
                    "max_subtasks": context.max_subtasks,
                    "estimator_name": context.estimator_name,
                    "source_estimator_name": context.source_estimator_name,
                    "source_task_name": context.source_task_name,
                },
            )

    def on_fit_task_end(self, estimator, context):
        # The path is enough to update the progress of the task and its ancestors.
        send(
            self._listener_handles[context.root_uuid],
            {
                "event": "end",
                "path": [ctx.task_id for ctx in get_context_path(context)],
            },
        )

    def teardown(self, estimator, context):
        # Fit is finished. Signal that the queue won't receive any more tasks, close
        # the monitor thread and the listener.
        _run_queues.pop(context.root_uuid).put(None)
        _run_monitors.pop(context.root_uuid).join()
        close_listener(self._listener_handles.pop(context.root_uuid))


try:
    from rich.progress import BarColumn, Progress, TextColumn, TimeRemainingColumn
    from rich.style import Style
    from rich.text import Text

    class _Progress(Progress):
        # Custom Progress class to allow showing the tasks in a given order (given by
        # setting the _ordered_tasks attribute). In particular it allows to dynamically
        # create and insert tasks between existing tasks.
        def get_renderables(self):
            yield self.make_tasks_table(getattr(self, "_ordered_tasks", []))

    class _StyledTimeRemainingColumn(TimeRemainingColumn):
        """Time column with millisecond precision and color styling."""

        def render(self, task):
            if self.elapsed_when_finished and task.finished:
                seconds = task.elapsed
            else:
                seconds = task.time_remaining

            if seconds is None:
                return Text("-:--:--.---")

            ms = round(seconds % 1 * 1000)
            text = Text(f"{timedelta(seconds=int(seconds))}.{ms:03d}")
            text.style = "#29ABE2" if task.finished else "#F7931E"
            return text

    class _StyledPercentageColumn(TextColumn):
        """Percentage column with color styling."""

        def render(self, task):
            text = super().render(task)
            text.style = "#29ABE2" if task.finished else "#F7931E"
            return text

except ImportError:
    pass


class RichProgressMonitor(Thread):
    """Thread monitoring the progress of an estimator with rich based display.

    The display is a list of nested rich tasks using `rich.Progress`. There is one for
    each non-leaf node in the task tree of the estimator.

    Parameters
    ----------
    queue : `queue.Queue` instance
        This thread will run until the queue is empty.
    """

    def __init__(self, *, queue):
        super().__init__()
        self.queue = queue

    def run(self):
        self.progress_ctx = _Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(
                complete_style=Style(color="#F7931E"),
                finished_style=Style(color="#29ABE2"),
                pulse_style=Style(color="#F7931E"),
            ),
            _StyledPercentageColumn("{task.percentage:>3.0f}%"),
            _StyledTimeRemainingColumn(elapsed_when_finished=True),
        )

        # Holds the root of the tree of rich tasks (i.e. progress bars) that will be
        # created dynamically as the task tree of the estimator is traversed.
        self.root_rich_task = None

        with self.progress_ctx:
            while task_info := self.queue.get():
                if task_info.pop("event") == "begin":
                    self._on_task_begin(task_info)
                else:
                    self._on_task_end(task_info)

            self.progress_ctx.refresh()

    def _on_task_begin(self, task_info):
        """Create a progress bar for the task and update the list of ordered tasks."""
        path = task_info.pop("path")

        rich_task = RichTask(
            progress_ctx=self.progress_ctx, task_info=task_info, depth=len(path) - 1
        )
        if rich_task.depth == 0:
            self.root_rich_task = rich_task
        else:
            parent = self.root_rich_task.get_descendants(path)[-2]
            parent.children[path[-1]] = rich_task

        self.progress_ctx._ordered_tasks = [
            self.progress_ctx.tasks[task.id] for task in self.root_rich_task
        ]

    def _on_task_end(self, task_info):
        """Update the progress of the task and its ancestors recursively."""
        path = task_info.pop("path")
        *ancestors, task = self.root_rich_task.get_descendants(path)

        if task is None:
            # a leaf task of the estimator, no progress bar was created for it
            task = RichTask(self.progress_ctx, task_info, depth=len(ancestors))
            ancestors[-1].children[path[-1]] = task
        else:
            # Task is finished. Render the progress bar as 100% completed regardless of
            # its progress because a task may execute less tasks than its max_subtasks.
            task.completed = task.total
            self.progress_ctx.update(task.id, completed=1, total=1)

        for ancestor in reversed(ancestors):
            if ancestor.total is None:
                continue

            completed = ancestor.progress * ancestor.total
            self.progress_ctx.update(ancestor.id, completed=completed)


class RichTask:
    """A task, i.e. progressbar, in the tree of rich tasks.

    There is a rich task for each non-leaf task in the context tree of the estimator.

    Parameters
    ----------
    progress_ctx : `rich.Progress` instance
        The progress context to which this task belongs.

    task_info : dict
        Information about the task for which this rich task is created.

    depth : int
        The depth of the task in the tree of rich tasks.

    Attributes
    ----------
    completed : int
        The number of completed subtasks.

    total : int or None
        The total number of subtasks. None if the total number of subtasks is not known.

    progress : float
        The fraction, between 0 and 1, of the task that is completed.

    id : int
        The ID of the task in the Progress context.

    children : dict
        A mapping from the index of a child to the child node `{idx: RichTask}`.
        For a leaf, it's an empty dictionary.
    """

    def __init__(self, progress_ctx, task_info, *, depth):
        self.children = {}
        self.depth = depth
        self.completed = 0
        self.total = task_info.pop("max_subtasks", 0)

        if task_info:
            description = self._format_task_description(task_info)
            self.id = progress_ctx.add_task(description, total=self.total)

    @property
    def progress(self):
        """Return the fraction of the task that is completed."""
        if self.completed == self.total:
            return 1.0
        if self.total is None:
            return 0.0
        return sum(child.progress for child in self.children.values()) / self.total

    def _format_task_description(self, task_info):
        """Return a formatted description for the task."""
        indent = f"{'  ' * self.depth}"
        task_desc = f"{task_info['estimator_name']} - {task_info['task_name']}"
        id_mark = f" #{task_info['task_id']}" if self.depth > 0 else ""
        source_task_desc = (
            f"{task_info['source_estimator_name']} - {task_info['source_task_name']} | "
            if task_info["source_estimator_name"] is not None
            else ""
        )
        return f"{indent}{source_task_desc}{task_desc}{id_mark}"

    def get_descendants(self, path):
        """Return the descendants from this task along the given path."""
        if len(path) == 1:
            return [self]
        if path[1] not in self.children:
            return [self, None]
        return [self] + self.children[path[1]].get_descendants(path[1:])

    def __iter__(self):
        """Pre-order depth-first traversal, only including tasks with a progress bar."""
        if hasattr(self, "id"):
            yield self
            for child in self.children.values():
                yield from child