299 lines
11 KiB
Python
299 lines
11 KiB
Python
# Authors: The scikit-learn developers
|
|
# SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
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
|
|
|
|
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 _StyledColumnMixin:
|
|
"""Apply finished/in-progress color style to rendered text."""
|
|
|
|
def render(self, task):
|
|
text = super().render(task)
|
|
text.style = "#29ABE2" if task.finished else "#F7931E"
|
|
return text
|
|
|
|
class _StyledTimeRemainingColumn(_StyledColumnMixin, TimeRemainingColumn):
|
|
"""Time column with color styling."""
|
|
|
|
class _StyledPercentageColumn(_StyledColumnMixin, TextColumn):
|
|
"""Percentage column with color styling."""
|
|
|
|
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
|