Dashboard

Nobody has claimed this yet.

### Description & Motivation

Lightning incorporates a number of 3rd party tools to visualize training (e.g., Tensorboard). However, during debugging or manual hyperparameter search, I mostly use console and would find it useful to be able to quickly judge whether the code is working as expected (typically whether the loss is going down; also whether model inputs and outputs appear reasonable).

### Pitch

I want to propose including a console-based dashboard that would incorporate user-defined components, such as progress bar, continuous plotting, or sample outputs from the model. Below is the code I've been using for myself for a while now that provides the described functionality. I attempts to build on top of `RichProgressBar`, so hopefully it wouldn't be too hard to issue a pull request if people get interested in this feature.

(For full context, see [my repo](https://github.com/qbilius/audible-papers) where this code is actually used for training.)

# `rich_progress.py`

Contains the base class for defining dashboards, `BaseDashboard`

```python
from threading import RLock
from typing import Any, Callable, Dict, Optional, Union

from lightning_utilities.core.imports import RequirementCache

import lightning.pytorch as pl
from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBar, RichProgressBarTheme, CustomProgress, MetricsTextColumn

GetTimeCallable = Callable[[], float]
_RICH_AVAILABLE = RequirementCache("rich>=10.2.2")

if _RICH_AVAILABLE:
from rich import get_console, reconfigure
from rich.console import Console
from rich.layout import Layout
from rich.live import Live
from rich.progress import ProgressColumn, Task, TaskID

class ProgressWithLayout(CustomProgress):
"""
Overrides ``CustomProgress`` to support passing layout as a renderable.

Presumably `get_renderable` would provide the same functionality without having to copy-paste the entire `__init__` method, but I cannot get it to work.
"""

def __init__(
self,
*columns: Union[str, ProgressColumn],
renderable=None, # new parameter
console: Optional[Console] = None,
auto_refresh: bool = True,
refresh_per_second: float = 10,
speed_estimate_period: float = 30.0,
transient: bool = False,
redirect_stdout: bool = True,
redirect_stderr: bool = True,
get_time: Optional[GetTimeCallable] = None,
disable: bool = False,
expand: bool = False,
) -> None:
assert refresh_per_second > 0, "refresh_per_second must be > 0"
self._lock = RLock()
self.columns = columns or self.get_default_columns()
self.speed_estimate_period = speed_estimate_period

self.disable = disable
self.expand = expand
self._tasks: Dict[TaskID, Task] = {}
self._task_index: TaskID = TaskID(0)
self.live = Live(
renderable=renderable, # pass it here
console=console or get_console(),
auto_refresh=auto_refresh,
refresh_per_second=refresh_per_second,
transient=transient,
redirect_stdout=redirect_stdout,
redirect_stderr=redirect_stderr,
)
self.get_time = get_time or self.console.get_time
self.print = self.console.print
self.log = self.console.log

class BaseDashboard(RichProgressBar):

"""Create a dashboard with a progress bar with `rich text formatting `_ and any other rich components.

This is a base class that should be inherited when composing your own dashboard. Your own subclass must define `layout` (`rich.layout.Layout`) that is split into subpanels (with unique names) and a `components` dictionary that for each of these names defines a subclass of `rich.jupyter.JupyterMixin` where the component's behavior is defined.

Args:
refresh_rate: Determines at which rate (in number of batches) the progress bars get updated.
Set it to ``0`` to disable the display.
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 = 1,
leave: bool = False,
theme: RichProgressBarTheme = RichProgressBarTheme(),
console_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
super().__init__(refresh_rate=refresh_rate,
leave=leave,
theme=theme,
console_kwargs=console_kwargs
)
self.layout: Optional[Layout] = None
self.components: Dict[str, Any] = {}

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()
self._console.clear_live()
self._metric_component = MetricsTextColumn(
trainer,
self.theme.metrics,
self.theme.metrics_text_delimiter,
self.theme.metrics_format,
)
self.progress = ProgressWithLayout(
*self.configure_columns(trainer),
self._metric_component,
auto_refresh=False,
disable=self.is_disabled,
console=self._console,
renderable=self.layout
)
self.progress.start()
# progress has started
self._progress_stopped = False

def _update_metrics(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
metrics = self.get_metrics(trainer, pl_module)
if self._metric_component:
self._metric_component.update(metrics)

for component in self.components.values():
component.update(trainer)
```

# `dashboard.py`

Contains a used-defined `Dashboard` with a progress bar, loss plot, and text samples output. `Dashboard` is passed to `Trainer` as any other callable.

```python
from typing import Any, Dict, Optional
from collections import deque, defaultdict

from lightning.pytorch.callbacks.progress.rich_progress import RichProgressBarTheme, _detect_light_colab_theme

from rich.ansi import AnsiDecoder
from rich.console import Group
from rich.jupyter import JupyterMixin
from rich.layout import Layout
from rich.progress import TaskID
from rich.table import Table

import plotext as plt

import rich_progress

class Dashboard(rich_progress.BaseDashboard):

def __init__(self,
plot_window: Optional[int] = None,
refresh_rate: int = 1,
leave: bool = False,
theme: RichProgressBarTheme = RichProgressBarTheme(),
console_kwargs: Optional[Dict[str, Any]] = None,
):
super().__init__(refresh_rate=refresh_rate,
leave=leave,
theme=theme,
console_kwargs=console_kwargs
)

self.components = {
'plot': Plot(window=plot_window),
'samples': TextSamples()
}
self.layout: Layout = Layout(name='root')
self.layout.split(
Layout(name='plot', ratio=1),
Layout(name='progress_bar', size=1),
Layout(name='samples', size=6),
)

def _update(self,
progress_bar_id: Optional['TaskID'],
current: int,
visible: bool = True
) -> None:

super()._update(
progress_bar_id=progress_bar_id,
current=current,
visible=visible
)

# make room for multiple progress bars
self.layout['progress_bar'].size = sum([t.visible for t in self.progress.tasks])
self.layout['progress_bar'].update(self.progress)

for name, component in self.components.items():
self.layout[name].update(component)

self.refresh()

class Plot(JupyterMixin):

def __init__(self, window: Optional[int] = None):
self.window = window

self._metrics = defaultdict(lambda: [
deque(maxlen=window), deque(maxlen=window)])
self.decoder = AnsiDecoder()

def __rich_console__(self, console, options):
self.width = options.max_width or console.width
self.height = options.height or console.height
canvas = self.make_plot()
self.rich_canvas = Group(*self.decoder.decode(canvas))
yield self.rich_canvas

def update(self, trainer) -> None:
for name, value in trainer.progress_bar_metrics.items():
self._metrics[name][0].append(trainer.global_step)
self._metrics[name][1].append(value)

def make_plot(self):
plt.clear_data()
plt.clear_figure()

xs = []
ys = []
for name, (x, y) in self._metrics.items():
plt.plot(x, y, label=name)
xs.extend(list(x))
if name.startswith('val'):
ys.extend(list(y))

if not _detect_light_colab_theme():
plt.theme('dark')
plt.plotsize(self.width, self.height)

if self.window is not None:
start = max(0, max(xs, default=0) - self.window)
plt.xlim(start, start + self.window)

return plt.build()

class TextSamples(JupyterMixin):

def __init__(self):
self.samples = []

def __rich_console__(self, console, options):
grid = Table.grid()
for sample in self.samples:
grid.add_row(sample)
yield grid

def update(self, trainer) -> None:
self.samples = trainer.loggers[0].samples
```

# `loggers.py`

Since dashboard is displaying sample output from the model, I additionally define a `TextLogger` class that stores the latest batch, so that it can be printed in the dashboard.

```python
from lightning.pytorch.loggers.logger import DummyLogger
from lightning.fabric.loggers.csv_logs import CSVLogger as FabricCSVLogger

from lightning.fabric.utilities.types import _PATH

class TextLogger(DummyLogger):

def __init__(self, save_dir: _PATH):
# A quick hack to the version number
self._version = FabricCSVLogger(root_dir=save_dir).version
self.samples = []

@property
def version(self):
return self._version

def log_samples(self, *samples):
self.samples = samples
```

### Alternatives

_No response_

### Additional context

This is how it looks like in practice. Train and validation losses are continuously updated and sample outputs are printed (in this case, I was training nanoGPT to remove citations from academic papers and you can see how the output does not contain `(Plyaskine et al., 1991)`).

cc @lantiga @borda

Start by reviewing the existing RichProgressBar integration and the proposed rich_progress.py, dashboard.py, and loggers.py examples. Clarify which dashboard components and Trainer/logger interfaces should be supported before implementation. Done should include an agreed console dashboard design that can display progress, metrics plots, and model samples.

Written by the indexing model from the issue text.