# Copyright The Lightning AI team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import time
from collections.abc import Generator
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, Optional, Union, cast
import torch
from lightning_utilities.core.apply_func import apply_to_collection
from typing_extensions import override
import lightning.pytorch as pl
from lightning.fabric.utilities.imports import _IS_INTERACTIVE
from lightning.pytorch.callbacks.progress.progress_bar import ProgressBar
from lightning.pytorch.utilities.imports import _RICH_AVAILABLE
from lightning.pytorch.utilities.types import STEP_OUTPUT
if _RICH_AVAILABLE:
from rich import get_console, reconfigure
from rich.console import Console, RenderableType
from rich.live import _RefreshThread as _RichRefreshThread
from rich.progress import BarColumn, Progress, ProgressColumn, Task, TaskID, TextColumn
from rich.progress_bar import ProgressBar as _RichProgressBar
from rich.style import Style
from rich.text import Text
class CustomBarColumn(BarColumn):
"""Overrides ``BarColumn`` to provide support for dataloaders that do not define a size (infinite size) such as
``IterableDataset``."""
def render(self, task: "Task") -> _RichProgressBar:
"""Gets a progress bar widget for a task."""
assert task.total is not None
assert task.remaining is not None
return _RichProgressBar(
total=max(0, task.total),
completed=max(0, task.completed),
width=None if self.bar_width is None else max(1, self.bar_width),
pulse=not task.started or not math.isfinite(task.remaining),
animation_time=task.get_time(),
style=self.style,
complete_style=self.complete_style,
finished_style=self.finished_style,
pulse_style=self.pulse_style,
)
@dataclass
class CustomInfiniteTask(Task):
"""Overrides ``Task`` to define an infinite task.
This is useful for datasets that do not define a size (infinite size) such as ``IterableDataset``.
"""
@property
def time_remaining(self) -> Optional[float]:
return None
class _RefreshThread(_RichRefreshThread):
def __init__(self, *args: Any, **kwargs: Any) -> None:
self.refresh_cond = False
super().__init__(*args, **kwargs)
def run(self) -> None:
while not self.done.is_set():
if self.refresh_cond:
with self.live._lock:
self.live.refresh()
self.refresh_cond = False
time.sleep(1 / self.refresh_per_second)
class CustomProgress(Progress):
"""Overrides ``Progress`` to support adding tasks that have an infinite total size."""
def start(self) -> None:
"""Starts the progress display.
Notes
-----
This override is needed to support the custom refresh thread.
"""
if self.live.auto_refresh:
self.live._refresh_thread = _RefreshThread(self.live, self.live.refresh_per_second)
self.live.auto_refresh = False
super().start()
if self.live._refresh_thread:
self.live.auto_refresh = True
self.live._refresh_thread.start()
def stop(self) -> None:
refresh_thread = self.live._refresh_thread
super().stop()
if refresh_thread:
refresh_thread.stop()
refresh_thread.join()
def soft_refresh(self) -> None:
if self.live.auto_refresh and isinstance(self.live._refresh_thread, _RefreshThread):
self.live._refresh_thread.refresh_cond = True
def add_task(
self,
description: str,
start: bool = True,
total: Optional[float] = 100.0,
completed: int = 0,
visible: bool = True,
**fields: Any,
) -> TaskID:
assert total is not None
if not math.isfinite(total):
task = CustomInfiniteTask(
self._task_index,
description,
total,
completed,
visible=visible,
fields=fields,
_get_time=self.get_time,
_lock=self._lock,
)
return self.add_custom_task(task)
return super().add_task(description, start, total, completed, visible, **fields)
def add_custom_task(self, task: CustomInfiniteTask, start: bool = True) -> TaskID:
with self._lock:
self._tasks[self._task_index] = task
if start:
self.start_task(self._task_index)
new_task_index = self._task_index
self._task_index = TaskID(int(self._task_index) + 1)
self.refresh()
return new_task_index
class CustomTimeColumn(ProgressColumn):
# Only refresh twice a second to prevent jitter
max_refresh = 0.5
def __init__(self, style: Union[str, Style]) -> None:
self.style = style
super().__init__()
def render(self, task: "Task") -> Text:
elapsed = task.finished_time if task.finished else task.elapsed
remaining = task.time_remaining
elapsed_delta = "-:--:--" if elapsed is None else str(timedelta(seconds=int(elapsed)))
remaining_delta = "-:--:--" if remaining is None else str(timedelta(seconds=int(remaining)))
return Text(f"{elapsed_delta} • {remaining_delta}", style=self.style)
class BatchesProcessedColumn(ProgressColumn):
def __init__(self, style: Union[str, Style]):
self.style = style
super().__init__()
def render(self, task: "Task") -> RenderableType:
total = task.total if task.total != float("inf") else "--"
return Text(f"{int(task.completed)}/{total}", style=self.style)
class ProcessingSpeedColumn(ProgressColumn):
def __init__(self, style: Union[str, Style]):
self.style = style
super().__init__()
def render(self, task: "Task") -> RenderableType:
task_speed = f"{task.speed:>.2f}" if task.speed is not None else "0.00"
return Text(f"{task_speed}it/s", style=self.style)
class MetricsTextColumn(ProgressColumn):
"""A column containing text."""
def __init__(
self,
trainer: "pl.Trainer",
style: Union[str, "Style"],
text_delimiter: str,
metrics_format: str,
):
self._trainer = trainer
self._tasks: dict[Union[int, TaskID], Any] = {}
self._current_task_id = 0
self._metrics: dict[Union[str, Style], Any] = {}
self._style = style
self._text_delimiter = text_delimiter
self._metrics_format = metrics_format
super().__init__()
def update(self, metrics: dict[Any, Any]) -> None:
# Called when metrics are ready to be rendered.
# This is to prevent render from causing deadlock issues by requesting metrics
# in separate threads.
self._metrics = metrics
def render(self, task: "Task") -> Text:
assert isinstance(self._trainer.progress_bar_callback, RichProgressBar)
if (
self._trainer.state.fn != "fit"
or self._trainer.sanity_checking
or self._trainer.progress_bar_callback.train_progress_bar_id != task.id
):
return Text()
if self._trainer.training and task.id not in self._tasks:
self._tasks[task.id] = "None"
if self._renderable_cache and self._current_task_id in self._renderable_cache:
self._current_task_id = cast(TaskID, self._current_task_id)
self._tasks[self._current_task_id] = self._renderable_cache[self._current_task_id][1]
self._current_task_id = task.id
if self._trainer.training and task.id != self._current_task_id:
return self._tasks[task.id]
metrics_texts = self._generate_metrics_texts()
text = self._text_delimiter.join(metrics_texts)
return Text(text, justify="left", style=self._style)
def _generate_metrics_texts(self) -> Generator[str, None, None]:
for name, value in self._metrics.items():
if not isinstance(value, str):
try:
value = f"{value:{self._metrics_format}}"
except (TypeError, ValueError):
value = str(value)
yield f"{name}: {value}"
@dataclass
class RichProgressBarTheme:
"""Styles to associate to different base components.
Args:
description: Style for the progress bar description. For eg., Epoch x, Testing, etc.
progress_bar: Style for the bar in progress.
progress_bar_finished: Style for the finished progress bar.
progress_bar_pulse: Style for the progress bar when `IterableDataset` is being processed.
batch_progress: Style for the progress tracker (i.e 10/50 batches completed).
time: Style for the processed time and estimate time remaining.
processing_speed: Style for the speed of the batches being processed.
metrics: Style for the metrics
https://rich.readthedocs.io/en/stable/style.html
"""
description: Union[str, "Style"] = ""
progress_bar: Union[str, "Style"] = "#6206E0"
progress_bar_finished: Union[str, "Style"] = "#6206E0"
progress_bar_pulse: Union[str, "Style"] = "#6206E0"
batch_progress: Union[str, "Style"] = ""
time: Union[str, "Style"] = "dim"
processing_speed: Union[str, "Style"] = "dim underline"
metrics: Union[str, "Style"] = "italic"
metrics_text_delimiter: str = " "
metrics_format: str = ".3f"
[docs]class RichProgressBar(ProgressBar):
"""Create a progress bar with `rich text formatting <https://github.com/Textualize/rich>`_.
Install it with pip:
.. code-block:: bash
pip install rich
.. code-block:: python
from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import RichProgressBar
trainer = Trainer(callbacks=RichProgressBar())
Args:
refresh_rate: Determines at which rate (per second) the progress bars get updated.
Set it to ``0`` to disable the display. Default: 100
leave: Leaves the finished progress bar in the terminal at the end of the epoch. Default: False
theme: Contains styles used to stylize the progress bar.
console_kwargs: Args for constructing a `Console`
Raises:
ModuleNotFoundError:
If required `rich` package is not installed on the device.
Note:
PyCharm users will need to enable “emulate terminal” in output console option in
run/debug configuration to see styled output.
Reference: https://rich.readthedocs.io/en/latest/introduction.html#requirements
"""
def __init__(
self,
refresh_rate: int = 100,
leave: bool = False,
theme: RichProgressBarTheme = RichProgressBarTheme(),
console_kwargs: Optional[dict[str, Any]] = None,
) -> None:
if not _RICH_AVAILABLE:
raise ModuleNotFoundError(
"`RichProgressBar` requires `rich` >= 10.2.2. Install it by running `pip install -U rich`."
)
super().__init__()
self._refresh_rate: int = refresh_rate
self._leave: bool = leave
self._console: Optional[Console] = None
self._console_kwargs = console_kwargs or {}
self._enabled: bool = True
self.progress: Optional[CustomProgress] = None
self.train_progress_bar_id: Optional[TaskID]
self.val_sanity_progress_bar_id: Optional[TaskID] = None
self.val_progress_bar_id: Optional[TaskID]
self.test_progress_bar_id: Optional[TaskID]
self.predict_progress_bar_id: Optional[TaskID]
self._reset_progress_bar_ids()
self._metric_component: Optional[MetricsTextColumn] = None
self._progress_stopped: bool = False
self.theme = theme
@property
def refresh_rate(self) -> float:
return self._refresh_rate
@property
def is_enabled(self) -> bool:
return self._enabled and self.refresh_rate > 0
@property
def is_disabled(self) -> bool:
return not self.is_enabled
@property
def train_progress_bar(self) -> "Task":
assert self.progress is not None
assert self.train_progress_bar_id is not None
return self.progress.tasks[self.train_progress_bar_id]
@property
def val_sanity_check_bar(self) -> "Task":
assert self.progress is not None
assert self.val_sanity_progress_bar_id is not None
return self.progress.tasks[self.val_sanity_progress_bar_id]
@property
def val_progress_bar(self) -> "Task":
assert self.progress is not None
assert self.val_progress_bar_id is not None
return self.progress.tasks[self.val_progress_bar_id]
@property
def test_progress_bar(self) -> "Task":
assert self.progress is not None
assert self.test_progress_bar_id is not None
return self.progress.tasks[self.test_progress_bar_id]
[docs] @override
def disable(self) -> None:
self._enabled = False
[docs] @override
def enable(self) -> None:
self._enabled = True
def _init_progress(self, trainer: "pl.Trainer") -> None:
if self.is_enabled and (self.progress is None or self._progress_stopped):
self._reset_progress_bar_ids()
reconfigure(**self._console_kwargs)
self._console = get_console()
# Compatibility shim for Rich >= 14.1.0:
if hasattr(self._console, "_live_stack"):
# In recent Rich releases, the internal `_live` variable was replaced with `_live_stack` (a list)
# to support nested Live displays. This broke our original call to `clear_live()`,
# because it now only pops one Live instance instead of clearing them all.
# We check for `_live_stack` and clear it manually for compatibility across
# both old and new Rich versions.
if len(self._console._live_stack) > 0:
self._console.clear_live()
else:
self._console.clear_live()
self._metric_component = MetricsTextColumn(
trainer,
self.theme.metrics,
self.theme.metrics_text_delimiter,
self.theme.metrics_format,
)
self.progress = CustomProgress(
*self.configure_columns(trainer),
self._metric_component,
auto_refresh=True,
refresh_per_second=self.refresh_rate if self.is_enabled else 1,
disable=self.is_disabled,
console=self._console,
)
self.progress.start()
# progress has started
self._progress_stopped = False
def refresh(self, hard: bool = False) -> None:
if self.progress:
if hard or _IS_INTERACTIVE:
self.progress.refresh()
else:
self.progress.soft_refresh()
[docs] @override
def on_train_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self._init_progress(trainer)
[docs] @override
def on_predict_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self._init_progress(trainer)
[docs] @override
def on_test_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self._init_progress(trainer)
[docs] @override
def on_validation_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self._init_progress(trainer)
[docs] @override
def on_sanity_check_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self._init_progress(trainer)
[docs] @override
def on_sanity_check_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
if self.progress is not None and self.val_sanity_progress_bar_id is not None:
self.progress.update(self.val_sanity_progress_bar_id, advance=0, visible=False)
self.refresh()
[docs] @override
def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
if self.is_disabled:
return
total_batches = self.total_train_batches
train_description = self._get_train_description(trainer.current_epoch)
if self.train_progress_bar_id is not None and self._leave:
self._stop_progress()
self._init_progress(trainer)
if self.progress is not None:
if self.train_progress_bar_id is None:
self.train_progress_bar_id = self._add_task(total_batches, train_description)
else:
self.progress.reset(
self.train_progress_bar_id,
total=total_batches,
description=train_description,
visible=True,
)
self.refresh()
[docs] @override
def on_validation_batch_start(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
batch: Any,
batch_idx: int,
dataloader_idx: int = 0,
) -> None:
if self.is_disabled or not self.has_dataloader_changed(dataloader_idx):
return
assert self.progress is not None
if trainer.sanity_checking:
if self.val_sanity_progress_bar_id is not None:
self.progress.update(self.val_sanity_progress_bar_id, advance=0, visible=False)
self.val_sanity_progress_bar_id = self._add_task(
self.total_val_batches_current_dataloader,
self.sanity_check_description,
visible=False,
)
else:
if self.val_progress_bar_id is not None:
self.progress.update(self.val_progress_bar_id, advance=0, visible=False)
# TODO: remove old tasks when new once they are created
self.val_progress_bar_id = self._add_task(
self.total_val_batches_current_dataloader,
self.validation_description,
visible=False,
)
self.refresh()
def _add_task(self, total_batches: Union[int, float], description: str, visible: bool = True) -> "TaskID":
assert self.progress is not None
return self.progress.add_task(
f"[{self.theme.description}]{description}" if self.theme.description else description,
total=total_batches,
visible=visible,
)
def _initialize_train_progress_bar_id(self) -> None:
total_batches = self.total_train_batches
train_description = self._get_train_description(self.trainer.current_epoch)
self.train_progress_bar_id = self._add_task(total_batches, train_description)
def _update(
self,
progress_bar_id: Optional["TaskID"],
current: int,
visible: bool = True,
hard: bool = False,
) -> None:
if self.progress is not None and self.is_enabled and progress_bar_id is not None:
self.progress.update(progress_bar_id, completed=current, visible=visible)
self.refresh(hard=hard)
[docs] @override
def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
if self.is_enabled and self.val_progress_bar_id is not None and trainer.state.fn == "fit":
assert self.progress is not None
self.progress.update(self.val_progress_bar_id, advance=0, visible=False)
self.refresh()
[docs] @override
def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
if trainer.state.fn == "fit":
self._update_metrics(trainer, pl_module)
self.reset_dataloader_idx_tracker()
[docs] @override
def on_test_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self.reset_dataloader_idx_tracker()
[docs] @override
def on_predict_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self.reset_dataloader_idx_tracker()
[docs] @override
def on_test_batch_start(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
batch: Any,
batch_idx: int,
dataloader_idx: int = 0,
) -> None:
if self.is_disabled or not self.has_dataloader_changed(dataloader_idx):
return
if self.test_progress_bar_id is not None:
assert self.progress is not None
self.progress.update(self.test_progress_bar_id, advance=0, visible=False)
self.test_progress_bar_id = self._add_task(self.total_test_batches_current_dataloader, self.test_description)
self.refresh()
[docs] @override
def on_predict_batch_start(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
batch: Any,
batch_idx: int,
dataloader_idx: int = 0,
) -> None:
if self.is_disabled or not self.has_dataloader_changed(dataloader_idx):
return
if self.predict_progress_bar_id is not None:
assert self.progress is not None
self.progress.update(self.predict_progress_bar_id, advance=0, visible=False)
self.predict_progress_bar_id = self._add_task(
self.total_predict_batches_current_dataloader, self.predict_description
)
self.refresh()
[docs] @override
def on_train_batch_end(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
outputs: STEP_OUTPUT,
batch: Any,
batch_idx: int,
) -> None:
if not self.is_disabled and self.train_progress_bar_id is None:
# can happen when resuming from a mid-epoch restart
self._initialize_train_progress_bar_id()
self._update(self.train_progress_bar_id, batch_idx + 1)
self._update_metrics(trainer, pl_module)
self.refresh()
[docs] @override
def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
self._update_metrics(trainer, pl_module)
self.refresh()
[docs] @override
def on_validation_batch_end(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
outputs: STEP_OUTPUT,
batch: Any,
batch_idx: int,
dataloader_idx: int = 0,
) -> None:
if self.is_disabled:
return
if trainer.sanity_checking:
if self.val_sanity_progress_bar_id is not None:
self._update(self.val_sanity_progress_bar_id, batch_idx + 1)
return
if self.val_progress_bar_id is None:
return
self._update(self.val_progress_bar_id, batch_idx + 1)
[docs] @override
def on_test_batch_end(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
outputs: STEP_OUTPUT,
batch: Any,
batch_idx: int,
dataloader_idx: int = 0,
) -> None:
if self.is_disabled or self.test_progress_bar_id is None:
return
self._update(self.test_progress_bar_id, batch_idx + 1)
[docs] @override
def on_predict_batch_end(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
outputs: Any,
batch: Any,
batch_idx: int,
dataloader_idx: int = 0,
) -> None:
if self.is_disabled or self.predict_progress_bar_id is None:
return
self._update(self.predict_progress_bar_id, batch_idx + 1)
def _get_train_description(self, current_epoch: int) -> str:
train_description = f"Epoch {current_epoch}"
if self.trainer.max_epochs is not None:
train_description += f"/{self.trainer.max_epochs - 1}"
if len(self.validation_description) > len(train_description):
# Padding is required to avoid flickering due of uneven lengths of "Epoch X"
# and "Validation" Bar description
train_description = f"{train_description:{len(self.validation_description)}}"
return train_description
def _stop_progress(self) -> None:
if self.progress is not None:
self.progress.stop()
# # signals for progress to be re-initialized for next stages
self._progress_stopped = True
def _reset_progress_bar_ids(self) -> None:
self.train_progress_bar_id = None
self.val_sanity_progress_bar_id = None
self.val_progress_bar_id = None
self.test_progress_bar_id = None
self.predict_progress_bar_id = None
[docs] @override
def get_metrics(
self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
) -> dict[str, Union[int, str, float, dict[str, float]]]:
items = super().get_metrics(trainer, pl_module)
# convert all metrics to float before sending to rich
return apply_to_collection(items, torch.Tensor, lambda x: x.item())
def _update_metrics(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
) -> None:
if not self.is_enabled or self._metric_component is None:
return
metrics = self.get_metrics(trainer, pl_module)
if self._metric_component:
self._metric_component.update(metrics)
[docs] @override
def teardown(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
self._stop_progress()
[docs] @override
def on_exception(
self,
trainer: "pl.Trainer",
pl_module: "pl.LightningModule",
exception: BaseException,
) -> None:
self._stop_progress()
def configure_columns(self, trainer: "pl.Trainer") -> list:
return [
TextColumn("[progress.description]{task.description}"),
CustomBarColumn(
complete_style=self.theme.progress_bar,
finished_style=self.theme.progress_bar_finished,
pulse_style=self.theme.progress_bar_pulse,
),
BatchesProcessedColumn(style=self.theme.batch_progress),
CustomTimeColumn(style=self.theme.time),
ProcessingSpeedColumn(style=self.theme.processing_speed),
]
def __getstate__(self) -> dict:
state = self.__dict__.copy()
# both the console and progress object can hold thread lock objects that are not pickleable
state["progress"] = None
state["_console"] = None
return state
