trainpit 0.1.0__py3-none-any.whl

trainpit/__init__.py

@@ -0,0 +1,5 @@
+"""Rich CLI progress monitoring for machine learning training loops."""
+
+from trainpit._tracker import TrainTracker, train
+
+__all__ = ["TrainTracker", "train"]

trainpit/_state.py

@@ -0,0 +1,195 @@
+"""Internal state for train progress tracking."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from time import monotonic
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+Scalar = float | int
+DisplayText = Annotated[str, Field(strict=True, min_length=1)]
+FiniteFloat = Annotated[float, Field(strict=True, allow_inf_nan=False)]
+NonNegativeFiniteFloat = Annotated[
+    float,
+    Field(strict=True, ge=0, allow_inf_nan=False),
+]
+PositiveInt = Annotated[int, Field(strict=True, ge=1)]
+TrainPhase = Literal["train"]
+
+
+class TrainState(BaseModel):
+    """Current state for one train display."""
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        extra="forbid",
+        strict=True,
+        validate_assignment=True,
+        validate_default=True,
+    )
+
+    label: DisplayText | None = Field(
+        default=None,
+        description="Human-readable label for the training run.",
+    )
+    total_epochs: PositiveInt | None = Field(
+        default=None,
+        description="Total epoch count when known.",
+    )
+    total_steps: PositiveInt | None = Field(
+        default=None,
+        description="Total step count per epoch when known.",
+    )
+    phase: TrainPhase = Field(
+        default="train",
+        description="Current lifecycle phase for the tracker.",
+    )
+    current_epoch: PositiveInt | None = Field(
+        default=None,
+        description="Current one-based epoch index.",
+    )
+    current_step: PositiveInt | None = Field(
+        default=None,
+        description="Current one-based step index within the active epoch.",
+    )
+    loss: FiniteFloat | None = Field(
+        default=None,
+        description="Latest finite loss value.",
+    )
+    metrics: dict[DisplayText, FiniteFloat] = Field(
+        default_factory=dict,
+        description="Latest finite scalar metrics keyed by metric name.",
+    )
+    learning_rate: NonNegativeFiniteFloat | None = Field(
+        default=None,
+        description="Latest non-negative learning rate.",
+    )
+    event: DisplayText | None = Field(
+        default=None,
+        description="Latest event message.",
+    )
+    started_at: NonNegativeFiniteFloat | None = Field(
+        default=None,
+        description="Monotonic timestamp when progress tracking started.",
+    )
+    updated_at: NonNegativeFiniteFloat | None = Field(
+        default=None,
+        description="Monotonic timestamp for the latest state update.",
+    )
+    finished_at: NonNegativeFiniteFloat | None = Field(
+        default=None,
+        description="Monotonic timestamp when progress tracking ended.",
+    )
+    started: bool = Field(
+        default=False,
+        description="Whether progress tracking has started.",
+    )
+    finished: bool = Field(
+        default=False,
+        description="Whether progress tracking has reached a terminal state.",
+    )
+    failed: bool = Field(
+        default=False,
+        description="Whether progress tracking ended with an error.",
+    )
+    error: BaseException | None = Field(
+        default=None,
+        description="Error captured when progress tracking fails.",
+    )
+
+    def start(self, *, now: Scalar | None = None) -> None:
+        timestamp = _resolve_timestamp(now)
+        self.started = True
+        self.finished = False
+        self.failed = False
+        self.error = None
+        self.started_at = timestamp
+        self.updated_at = timestamp
+        self.finished_at = None
+
+    def set_epoch(self, value: int, *, now: Scalar | None = None) -> None:
+        self.current_epoch = _require_positive_integer("epoch", value)
+        self._touch(_resolve_timestamp(now))
+
+    def set_step(
+        self,
+        value: int,
+        *,
+        loss: Scalar | None = None,
+        metrics: Mapping[str, Scalar] | None = None,
+        learning_rate: Scalar | None = None,
+        now: Scalar | None = None,
+    ) -> None:
+        self.current_step = _require_positive_integer("step", value)
+
+        if loss is not None:
+            self.loss = _coerce_scalar("loss", loss)
+        if metrics is not None:
+            self.metrics = {**self.metrics, **_coerce_metrics(metrics)}
+        if learning_rate is not None:
+            self.learning_rate = _coerce_scalar("learning_rate", learning_rate)
+
+        self._touch(_resolve_timestamp(now))
+
+    def set_metrics(
+        self,
+        values: Mapping[str, Scalar],
+        *,
+        now: Scalar | None = None,
+    ) -> None:
+        self.metrics = {**self.metrics, **_coerce_metrics(values)}
+        self._touch(_resolve_timestamp(now))
+
+    def set_event(self, message: str, *, now: Scalar | None = None) -> None:
+        self.event = message
+        self._touch(_resolve_timestamp(now))
+
+    def finish(self, *, now: Scalar | None = None) -> None:
+        timestamp = _resolve_timestamp(now)
+        self._touch(timestamp)
+        self.finished = True
+        self.failed = False
+        self.error = None
+        self.finished_at = timestamp
+
+    def fail(self, error: BaseException, *, now: Scalar | None = None) -> None:
+        timestamp = _resolve_timestamp(now)
+        self._touch(timestamp)
+        self.finished = True
+        self.failed = True
+        self.error = error
+        self.finished_at = timestamp
+
+    def _touch(self, timestamp: float) -> None:
+        if self.started_at is None:
+            self.started_at = timestamp
+        self.updated_at = timestamp
+
+
+def _require_positive_integer(name: str, value: int) -> int:
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise TypeError(f"{name} must be an integer")
+    if value < 1:
+        raise ValueError(f"{name} must be greater than or equal to 1")
+    return value
+
+
+def _coerce_scalar(name: str, value: Scalar) -> float:
+    if isinstance(value, bool) or not isinstance(value, int | float):
+        raise TypeError(f"{name} must be a finite scalar")
+
+    return float(value)
+
+
+def _resolve_timestamp(value: Scalar | None) -> float:
+    timestamp = monotonic() if value is None else _coerce_scalar("now", value)
+    if timestamp < 0:
+        raise ValueError("now must be greater than or equal to 0")
+
+    return timestamp
+
+
+def _coerce_metrics(values: Mapping[str, Scalar]) -> dict[str, float]:
+    return {name: _coerce_scalar(name, value) for name, value in values.items()}

trainpit/_tracker.py

@@ -0,0 +1,126 @@
+"""Public train tracker entry point implementation."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from types import TracebackType
+from typing import Self
+
+from trainpit._state import Scalar, TrainState
+
+
+class TrainTracker:
+    """Context manager used to update one train progress display."""
+
+    def __init__(
+        self,
+        *,
+        total_epochs: int | None = None,
+        total_steps: int | None = None,
+        label: str | None = None,
+    ) -> None:
+        self._state = TrainState(
+            label=label,
+            total_epochs=_optional_positive_integer("total_epochs", total_epochs),
+            total_steps=_optional_positive_integer("total_steps", total_steps),
+        )
+
+    @property
+    def snapshot(self) -> TrainState:
+        """Return a copy of the current progress state."""
+
+        return self._state.model_copy(deep=True)
+
+    def __enter__(self) -> Self:
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool:
+        del exc_type, traceback
+
+        if exc is None:
+            self.finish()
+        else:
+            self.fail(exc)
+
+        return False
+
+    def epoch(self, value: int) -> None:
+        self._state.set_epoch(value)
+
+    def start(self) -> None:
+        self._state.start()
+
+    def step(
+        self,
+        value: int,
+        *,
+        loss: Scalar | None = None,
+        metrics: Mapping[str, Scalar] | None = None,
+        learning_rate: Scalar | None = None,
+        lr: Scalar | None = None,
+    ) -> None:
+        self._state.set_step(
+            value,
+            loss=loss,
+            metrics=metrics,
+            learning_rate=_resolve_learning_rate(learning_rate, lr),
+        )
+
+    def update_metrics(self, values: Mapping[str, Scalar]) -> None:
+        self._state.set_metrics(values)
+
+    def metrics(self, values: Mapping[str, Scalar]) -> None:
+        self.update_metrics(values)
+
+    def log(self, message: str) -> None:
+        self._state.set_event(message)
+
+    def event(self, message: str) -> None:
+        self.log(message)
+
+    def finish(self) -> None:
+        self._state.finish()
+
+    def fail(self, error: BaseException) -> None:
+        self._state.fail(error)
+
+
+def train(
+    *,
+    total_epochs: int | None = None,
+    total_steps: int | None = None,
+    label: str | None = None,
+) -> TrainTracker:
+    """Create a tracker for one training loop."""
+
+    return TrainTracker(
+        total_epochs=total_epochs,
+        total_steps=total_steps,
+        label=label,
+    )
+
+
+def _optional_positive_integer(name: str, value: int | None) -> int | None:
+    if value is None:
+        return None
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise TypeError(f"{name} must be an integer")
+    if value < 1:
+        raise ValueError(f"{name} must be greater than or equal to 1")
+    return value
+
+
+def _resolve_learning_rate(
+    learning_rate: Scalar | None,
+    lr: Scalar | None,
+) -> Scalar | None:
+    if learning_rate is not None and lr is not None:
+        raise ValueError("Use either learning_rate or lr, not both")
+
+    return learning_rate if learning_rate is not None else lr

trainpit/_tui/__init__.py

@@ -0,0 +1,43 @@
+"""Internal modules backing the public trainpit.tui facade."""
+
+from __future__ import annotations
+
+from trainpit._tui.app import TrainDashboardApp as TrainDashboardApp
+from trainpit._tui.formatting import (
+    _format_curves as _format_curves,
+    _format_metrics as _format_metrics,
+    _format_progress as _format_progress,
+    _format_timing as _format_timing,
+)
+from trainpit._tui.graphs import line_graph as line_graph
+from trainpit._tui.graphs import scatter_graph as scatter_graph
+from trainpit._tui.panels import DashboardPanel as DashboardPanel
+from trainpit._tui.snapshot import TrainDashboardSnapshot as TrainDashboardSnapshot
+from trainpit._tui.types import CurveGraphRenderer as CurveGraphRenderer
+from trainpit._tui.types import DisplayText as DisplayText
+from trainpit._tui.types import FiniteFloat as FiniteFloat
+from trainpit._tui.types import NonNegativeFiniteFloat as NonNegativeFiniteFloat
+from trainpit._tui.types import PanelSlot as PanelSlot
+from trainpit._tui.types import PositiveInt as PositiveInt
+from trainpit._tui.types import RunStatus as RunStatus
+from trainpit._tui.types import Scalar as Scalar
+
+__all__ = [
+    "CurveGraphRenderer",
+    "DashboardPanel",
+    "DisplayText",
+    "FiniteFloat",
+    "NonNegativeFiniteFloat",
+    "PanelSlot",
+    "PositiveInt",
+    "RunStatus",
+    "Scalar",
+    "TrainDashboardApp",
+    "TrainDashboardSnapshot",
+    "_format_curves",
+    "_format_metrics",
+    "_format_progress",
+    "_format_timing",
+    "line_graph",
+    "scatter_graph",
+]

trainpit/_tui/app.py

@@ -0,0 +1,270 @@
+"""Textual app implementation for the trainpit dashboard."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+from textual.app import App, ComposeResult
+from textual.containers import Horizontal, Vertical
+from textual.widgets import Footer, Header, Static
+
+from trainpit._tui.formatting import (
+    _format_curves,
+    _format_metrics,
+    _format_progress,
+    _format_timing,
+)
+from trainpit._tui.graphs import line_graph
+from trainpit._tui.panels import (
+    DashboardPanel,
+    _panel_body_id,
+    _panel_container_id,
+    _panel_title,
+    _validate_extra_panels,
+)
+from trainpit._tui.snapshot import TrainDashboardSnapshot
+from trainpit._tui.types import CurveGraphRenderer, PanelSlot
+
+
+class TrainDashboardApp(App[None]):
+    """A compact Textual dashboard for train progress."""
+
+    CSS = """
+    Screen {
+        background: #07090d;
+        color: #d7dee8;
+    }
+
+    Header,
+    Footer {
+        background: #0b0f14;
+        color: #97a3b6;
+    }
+
+    #dashboard {
+        height: 1fr;
+        padding: 0 1 1 1;
+    }
+
+    #top-row {
+        height: 10;
+        min-height: 9;
+    }
+
+    #bottom-row {
+        height: 1fr;
+        min-height: 16;
+    }
+
+    .panel {
+        background: #0c1117;
+        border: round #27313f;
+        padding: 0 1;
+        margin: 0 1 1 0;
+        min-width: 24;
+    }
+
+    #status-panel {
+        width: 1fr;
+        min-width: 22;
+    }
+
+    #progress-panel {
+        width: 2fr;
+        min-width: 34;
+    }
+
+    #metrics-panel {
+        width: 2fr;
+        min-width: 32;
+    }
+
+    .graph-panel {
+        width: 4fr;
+        min-width: 60;
+    }
+
+    #bottom-side {
+        width: 1fr;
+        min-width: 24;
+    }
+
+    .side-panel {
+        height: 1fr;
+    }
+
+    .bottom-panel {
+        width: 2fr;
+        min-width: 32;
+    }
+
+    .custom-panel {
+        color: #d7dee8;
+    }
+
+    .panel-title {
+        background: #101b24;
+        color: #6ee7f9;
+        text-style: bold reverse;
+        padding: 0 1;
+        margin-bottom: 1;
+        width: 1fr;
+    }
+
+    #status {
+        color: #7ee787;
+        text-style: bold;
+    }
+
+    .status-running {
+        color: #7ee787;
+    }
+
+    .status-finished {
+        color: #6ee7f9;
+    }
+
+    .status-failed {
+        color: #ff7b72;
+    }
+
+    #label {
+        color: #f0f4fa;
+    }
+
+    #curve {
+        color: #7ee787;
+    }
+
+    #progress {
+        color: #f0c36a;
+    }
+
+    #metrics {
+        color: #f0f4fa;
+    }
+
+    #timing {
+        color: #97a3b6;
+    }
+
+    #events {
+        color: #d7dee8;
+    }
+    """
+
+    TITLE = "trainpit"
+    SUB_TITLE = "training monitor"
+    ENABLE_COMMAND_PALETTE = False
+    BINDINGS = [("q", "quit", "Quit")]
+
+    def __init__(
+        self,
+        snapshot: TrainDashboardSnapshot | None = None,
+        *,
+        extra_panels: Sequence[DashboardPanel] | None = None,
+        graph_renderer: CurveGraphRenderer = line_graph,
+    ) -> None:
+        super().__init__()
+        self.snapshot = snapshot or TrainDashboardSnapshot()
+        self.extra_panels = _validate_extra_panels(extra_panels or ())
+        self.graph_renderer = _validate_graph_renderer(graph_renderer)
+
+    def compose(self) -> ComposeResult:
+        yield Header(show_clock=True)
+        with Vertical(id="dashboard"):
+            with Horizontal(id="top-row"):
+                with Vertical(id="status-panel", classes="panel"):
+                    yield Static(_panel_title("STATUS"), classes="panel-title")
+                    yield Static("", id="status")
+                    yield Static("", id="label")
+                with Vertical(id="progress-panel", classes="panel"):
+                    yield Static(_panel_title("PROGRESS"), classes="panel-title")
+                    yield Static("", id="progress")
+                with Vertical(id="metrics-panel", classes="panel"):
+                    yield Static(_panel_title("METRICS"), classes="panel-title")
+                    yield Static("", id="metrics")
+                for panel in self._panels_for_slot("top"):
+                    with Vertical(
+                        id=_panel_container_id(panel),
+                        classes="panel custom-panel",
+                    ):
+                        yield Static(
+                            _panel_title(panel.title),
+                            classes="panel-title",
+                        )
+                        yield Static("", id=_panel_body_id(panel))
+            with Horizontal(id="bottom-row"):
+                with Vertical(classes="panel graph-panel"):
+                    yield Static(_panel_title("LEARNING CURVE"), classes="panel-title")
+                    yield Static("", id="curve")
+                for panel in self._panels_for_slot("bottom"):
+                    with Vertical(
+                        id=_panel_container_id(panel),
+                        classes="panel bottom-panel custom-panel",
+                    ):
+                        yield Static(
+                            _panel_title(panel.title),
+                            classes="panel-title",
+                        )
+                        yield Static("", id=_panel_body_id(panel))
+                with Vertical(id="bottom-side"):
+                    with Vertical(classes="panel side-panel"):
+                        yield Static(_panel_title("TIMING"), classes="panel-title")
+                        yield Static("", id="timing")
+                    with Vertical(classes="panel side-panel"):
+                        yield Static(_panel_title("EVENTS"), classes="panel-title")
+                        yield Static("", id="events")
+                    for panel in self._panels_for_slot("side"):
+                        with Vertical(
+                            id=_panel_container_id(panel),
+                            classes="panel side-panel custom-panel",
+                        ):
+                            yield Static(
+                                _panel_title(panel.title),
+                                classes="panel-title",
+                            )
+                            yield Static("", id=_panel_body_id(panel))
+        yield Footer()
+
+    def on_mount(self) -> None:
+        self.refresh_dashboard()
+
+    def refresh_dashboard(self) -> None:
+        """Render the current snapshot into dashboard widgets."""
+
+        snapshot = self.snapshot
+        status = self.query_one("#status", Static)
+        status.update(snapshot.status.upper())
+        for value in ("running", "finished", "failed"):
+            status.set_class(snapshot.status == value, f"status-{value}")
+
+        self.query_one("#label", Static).update(snapshot.label or "unlabeled run")
+        self.query_one("#progress", Static).update(_format_progress(snapshot))
+        self.query_one("#metrics", Static).update(_format_metrics(snapshot))
+        self.query_one("#curve", Static).update(
+            _format_curves(snapshot, graph_renderer=self.graph_renderer)
+        )
+        self.query_one("#timing", Static).update(_format_timing(snapshot))
+        self.query_one("#events", Static).update(snapshot.event or "waiting for events")
+        for panel in self.extra_panels:
+            self.query_one(f"#{_panel_body_id(panel)}", Static).update(
+                panel.render(snapshot)
+            )
+
+    def update_snapshot(self, snapshot: TrainDashboardSnapshot) -> None:
+        """Replace the dashboard snapshot and refresh visible widgets."""
+
+        self.snapshot = snapshot
+        if self.is_mounted:
+            self.refresh_dashboard()
+
+    def _panels_for_slot(self, slot: PanelSlot) -> tuple[DashboardPanel, ...]:
+        return tuple(panel for panel in self.extra_panels if panel.slot == slot)
+
+
+def _validate_graph_renderer(renderer: CurveGraphRenderer) -> CurveGraphRenderer:
+    if not callable(renderer):
+        raise TypeError("graph_renderer must be callable")
+
+    return renderer

trainpit/_tui/constants.py

@@ -0,0 +1,26 @@
+"""Constants used by trainpit TUI rendering."""
+
+from __future__ import annotations
+
+GRAPH_WIDTH = 56
+PRIMARY_GRAPH_HEIGHT = 7
+SECONDARY_GRAPH_HEIGHT = 3
+PROGRESS_BAR_WIDTH = 24
+BRAILLE_BASE = 0x2800
+BRAILLE_COLUMNS = (
+    (0x01, 0x02, 0x04, 0x40),
+    (0x08, 0x10, 0x20, 0x80),
+)
+BRAILLE_ROWS_PER_CELL = 4
+SAMPLES_PER_BRAILLE = 2
+EPOCH_AXIS_MIN_LABEL_SPACING = 8
+EPOCH_AXIS_LABEL_PADDING = 1
+PANEL_ID_CHARACTERS = frozenset(
+    "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-"
+)
+FULLWIDTH_TITLE_TRANSLATION = str.maketrans(
+    {
+        " ": " ",
+        **{chr(value): chr(value + 0xFEE0) for value in range(ord("A"), ord("Z") + 1)},
+    }
+)