"""Run monitoring for TRACMASS-JAX integration loops.
Provides :class:`RunMonitor`, a context-manager that receives structured
events from :func:`~tracmass.integrate.time_dependent` and renders them via
one of three backends:
- ``"print"`` — simple ``print()`` output, reproducing the original behaviour.
- ``"log"`` — structured :mod:`logging` records (one per window + summary).
- ``"rich"`` — live terminal dashboard powered by :mod:`rich`.
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
if TYPE_CHECKING:
import pandas as pd
# ---------------------------------------------------------------------------
# Formatting helpers
# ---------------------------------------------------------------------------
[docs]
def sec_str(s: float) -> str:
"""Format a duration in seconds as a compact human-readable string.
Parameters
----------
s : float
Duration in seconds.
Returns
-------
str
E.g. ``"42ms"``, ``"3.7s"``, ``"2m05s"``.
"""
if s < 1.0:
return f"{s * 1000:.0f}ms"
if s < 60.0:
return f"{s:.1f}s"
m = int(s // 60)
return f"{m}m{s - m * 60:.0f}s"
[docs]
def dt_str(dt) -> str:
"""Format a :class:`pandas.Timedelta` as a compact string.
Parameters
----------
dt : pandas.Timedelta
Returns
-------
str
"""
total = dt.total_seconds()
dy = int(total / 86400)
hr = int(total / 3600 - dy * 24)
mn = int(total / 60 - 60 * (dy * 24 + hr))
sk = int(total - 60 * (dy * 24 * 60 + hr * 60 + mn))
out = ""
out += "" if dy == 0 else f"{dy:02}d "
out += "" if hr == 0 else f"{hr:02}:"
out += "" if mn == 0 else f"{mn:02}:"
out += f"{sk}"
return out
def _reason_str(reason_tally: dict[str, int]) -> str:
"""Format a reason-code tally as a compact string."""
return ", ".join(f"{k}:{v}" for k, v in reason_tally.items() if v)
# ---------------------------------------------------------------------------
# RunMonitor
# ---------------------------------------------------------------------------
[docs]
class RunMonitor:
"""Monitor for :func:`~tracmass.integrate.time_dependent` runs.
Receives structured events from the integration loop and renders them
via the chosen backend. Use as a context manager so the ``"rich"``
backend can properly start and stop its live display.
Parameters
----------
backend : {"print", "log", "rich"}
Output backend. ``"print"`` reproduces the original stdout output.
``"log"`` emits :mod:`logging` records. ``"rich"`` shows a live
terminal dashboard.
logger : logging.Logger, optional
Logger used by the ``"log"`` backend. Defaults to the
``"tracmass"`` logger.
max_table_rows : int, optional
Maximum number of window rows shown in the ``"rich"`` table at any
one time. Older rows scroll out of view. Default ``10``.
Examples
--------
>>> with tracmass.RunMonitor(backend="rich") as mon:
... result = tracmass.time_dependent(state, grid, fluxes,
... n_windows=50, monitor=mon)
>>> # Plain print (default when monitor is not supplied):
>>> result = tracmass.time_dependent(state, grid, fluxes, n_windows=50)
"""
[docs]
def __init__(
self,
backend: str = "print",
logger: logging.Logger | None = None,
max_table_rows: int = 10,
) -> None:
if backend not in ("print", "log", "rich"):
raise ValueError(
f"backend must be 'print', 'log', or 'rich'; got {backend!r}"
)
self._backend = backend
self._logger = logger or logging.getLogger("tracmass")
self._max_table_rows = max_table_rows
# Timing accumulators
self._t_io: float = 0.0
self._t_jax: float = 0.0
self._t_py: float = 0.0
self._wins_done: int = 0
# Cumulative stopped-particle accumulators
self._cumul_stopped: int = 0
self._cumul_reason_tally: dict[str, int] = {}
# Set by on_run_start
self._n_windows: int | None = None
self._n_particles: int | None = None
# Rich state
self._live = None
self._progress = None
self._task_id = None
self._rows: list[tuple] = [] # all window rows (unbounded)
self._summary: tuple | None = None # totals row added on run_end
# ------------------------------------------------------------------
# Context manager
# ------------------------------------------------------------------
def __enter__(self) -> RunMonitor:
return self
def __exit__(self, *args) -> None:
if self._backend == "rich" and self._live is not None:
self._live.stop()
self._live = None
# ------------------------------------------------------------------
# Public event API
# ------------------------------------------------------------------
[docs]
def on_run_start(self, n_windows: int, n_particles: int) -> None:
"""Called once before the integration loop begins.
Parameters
----------
n_windows : int
Total number of GCM windows to integrate.
n_particles : int
Initial number of particles.
"""
self._n_windows = n_windows
self._n_particles = n_particles
if self._backend == "log":
self._logger.info(
"Run started: n_windows=%d, n_particles=%d",
n_windows,
n_particles,
)
elif self._backend == "rich" and self._live is None:
self._start_rich()
[docs]
def on_window_start(self, win: int, n_active: int, dtm: pd.Timestamp) -> None:
"""Called at the start of each window, before I/O and advection.
Parameters
----------
win : int
Zero-based window index.
n_active : int
Number of active particles entering this window.
dtm : pandas.Timestamp
Timestamp of the current GCM time window.
"""
if self._backend == "log":
self._logger.debug(
"Window %d/%d starting: %d active, dtm=%s",
win + 1,
self._n_windows,
n_active,
dtm.date() if dtm is not None else "?",
)
elif self._backend == "rich" and self._progress is not None:
total = self._n_windows if self._n_windows is not None else win + 1
n_left = total - win
self._progress.update(
self._task_id,
description=f"[cyan]{win} done, {n_left} left[/] …",
)
self._live.refresh()
[docs]
def on_window_end(
self,
win: int,
dtm: pd.Timestamp,
n_active_before: int,
n_active_after: int,
n_stopped: int,
reason_tally: dict[str, int],
n_sub: int,
limit_hit: bool,
t_io: float,
t_jax: float,
t_py: float,
) -> None:
"""Called at the end of each window after advection and reactivation.
Parameters
----------
win : int
Zero-based window index.
dtm : pandas.Timestamp
Timestamp of the window just completed.
n_active_before : int
Active count at window entry.
n_active_after : int
Active count after advection + reactivation.
n_stopped : int
Particles permanently deactivated this window.
reason_tally : dict[str, int]
Counts per deactivation reason name.
n_sub : int
Number of sub-cycles executed.
limit_hit : bool
Whether the sub-cycle limit was reached.
t_io : float
Wall time spent in I/O (Dask/Zarr reads), seconds.
t_jax : float
Wall time spent in JAX advection (device_put + while_loop), seconds.
t_py : float
Wall time spent in Python overhead (reactivation, GC, flux swap), seconds.
"""
self._t_io += t_io
self._t_jax += t_jax
self._t_py += t_py
self._wins_done += 1
self._cumul_stopped += n_stopped
for k, v in reason_tally.items():
self._cumul_reason_tally[k] = self._cumul_reason_tally.get(k, 0) + v
t_win = t_io + t_jax + t_py
if self._backend == "print":
self._print_window(
win,
dtm,
n_active_after,
n_stopped,
reason_tally,
n_sub,
limit_hit,
t_io,
t_jax,
t_py,
t_win,
)
elif self._backend == "log":
self._log_window(
win,
dtm,
n_active_before,
n_active_after,
n_stopped,
reason_tally,
n_sub,
limit_hit,
t_io,
t_jax,
t_py,
)
elif self._backend == "rich":
self._rich_window(
win,
dtm,
n_active_after,
n_stopped,
reason_tally,
n_sub,
limit_hit,
t_io,
t_jax,
t_py,
t_win,
)
[docs]
def on_run_end(self) -> None:
"""Called once after the integration loop finishes."""
t_total = self._t_io + self._t_jax + self._t_py
if self._backend == "print":
self._print_summary(t_total)
elif self._backend == "log":
self._log_summary(t_total)
elif self._backend == "rich":
self._rich_summary(t_total)
if self._live is not None:
self._live.stop()
self._live = None
# ------------------------------------------------------------------
# Print backend
# ------------------------------------------------------------------
def _print_window(
self,
win,
dtm,
n_active_after,
n_stopped,
reason_tally,
n_sub,
limit_hit,
t_io,
t_jax,
t_py,
t_win,
):
n_windows = self._n_windows if self._n_windows is not None else "?"
limit_str = " [LIMIT]" if limit_hit else ""
stop_str = (
f" stopped:{n_stopped} ({_reason_str(reason_tally)})" if n_stopped else ""
)
date_str = dtm.date() if dtm is not None else "?"
print(
f"{win + 1:03}/{n_windows:03} {date_str} "
f"{n_active_after} active"
f"{stop_str}, "
f"#sub:{n_sub}{limit_str}, "
f"wall:{sec_str(t_win)} "
f"[io:{sec_str(t_io)} "
f"jax:{sec_str(t_jax)} "
f"py:{sec_str(t_py)}]"
)
def _print_summary(self, t_total: float):
if self._wins_done <= 1 or t_total <= 0:
return
pct = lambda t: f"{100 * t / t_total:.0f}%" # noqa: E731
print(
f"--- total wall:{sec_str(t_total)} "
f"[io:{sec_str(self._t_io)} {pct(self._t_io)} "
f"jax:{sec_str(self._t_jax)} {pct(self._t_jax)} "
f"py:{sec_str(self._t_py)} {pct(self._t_py)}]"
)
if self._cumul_stopped:
reason_str = _reason_str(self._cumul_reason_tally)
print(f"--- stopped:{self._cumul_stopped} ({reason_str})")
# ------------------------------------------------------------------
# Log backend
# ------------------------------------------------------------------
def _log_window(
self,
win,
dtm,
n_active_before,
n_active_after,
n_stopped,
reason_tally,
n_sub,
limit_hit,
t_io,
t_jax,
t_py,
):
t_win = t_io + t_jax + t_py
date_str = str(dtm.date()) if dtm is not None else "?"
extra = {
"win": win + 1,
"n_windows": self._n_windows,
"dtm": date_str,
"n_active_before": n_active_before,
"n_active_after": n_active_after,
"n_stopped": n_stopped,
"n_sub": n_sub,
"limit_hit": limit_hit,
"t_io": round(t_io, 4),
"t_jax": round(t_jax, 4),
"t_py": round(t_py, 4),
"t_win": round(t_win, 4),
**{f"stopped_{k}": v for k, v in reason_tally.items()},
}
self._logger.info(
"window %d/%s %s: %d→%d active, stopped:%d, #sub:%d%s, "
"wall:%.2fs [io:%.2fs jax:%.2fs py:%.2fs]",
win + 1,
self._n_windows if self._n_windows is not None else "?",
date_str,
n_active_before,
n_active_after,
n_stopped,
n_sub,
" LIMIT" if limit_hit else "",
t_win,
t_io,
t_jax,
t_py,
extra=extra,
)
def _log_summary(self, t_total: float):
pct = lambda t: 100 * t / max(t_total, 1e-9) # noqa: E731
extra = {
"cumul_stopped": self._cumul_stopped,
**{f"stopped_{k}": v for k, v in self._cumul_reason_tally.items()},
}
self._logger.info(
"Run finished: %d windows, stopped:%d (%s), total %.2fs "
"[io:%.2fs %.0f%% jax:%.2fs %.0f%% py:%.2fs %.0f%%]",
self._wins_done,
self._cumul_stopped,
_reason_str(self._cumul_reason_tally),
t_total,
self._t_io,
pct(self._t_io),
self._t_jax,
pct(self._t_jax),
self._t_py,
pct(self._t_py),
extra=extra,
)
# ------------------------------------------------------------------
# Rich backend
# ------------------------------------------------------------------
def _make_table(self):
"""Build a fresh Table from the last ``max_table_rows`` stored rows."""
from rich.table import Table
from tracmass.config import settings
# Read on every rebuild so an env switched mid-run (``setenv``) shows up.
env = str(settings.current_env).lower()
table = Table(
"Date",
"Active",
"Stopped ∑",
"#sub",
"wall",
"io",
"jax",
"py",
title=f"tracmass · time_dependent · env: {env}",
show_lines=False,
header_style="bold cyan",
border_style="dim",
)
for row in self._rows[-self._max_table_rows :]:
table.add_row(*row)
if self._summary is not None:
table.add_section()
table.add_row(*self._summary)
return table
def _start_rich(self):
from rich.console import Group
from rich.live import Live
from rich.progress import (
BarColumn,
MofNCompleteColumn,
Progress,
SpinnerColumn,
TextColumn,
TimeElapsedColumn,
)
# `or 1` would wrongly fall back to 1 for a genuine n_windows=0 (e.g.
# a caller's off-by-typo date range yielding zero windows), leaving
# the bar frozen at a misleading "0/1" forever since nothing will
# ever advance it. Only an unset (None) total falls back to 1.
n_total = self._n_windows if self._n_windows is not None else 1
self._progress = Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
MofNCompleteColumn(),
TimeElapsedColumn(),
)
self._task_id = self._progress.add_task("", total=n_total)
self._live = Live(
Group(self._make_table(), self._progress),
refresh_per_second=4,
)
self._live.start()
def _rich_window(
self,
win,
dtm,
n_active_after,
n_stopped,
reason_tally,
n_sub,
limit_hit,
t_io,
t_jax,
t_py,
t_win,
):
if self._live is None:
return
from rich.console import Group
limit_str = " !" if limit_hit else ""
cumul = self._cumul_stopped
date_str = str(dtm.date()) if dtm is not None else "?"
self._rows.append(
(
date_str,
str(n_active_after),
f"[red]{cumul}[/]" if cumul else "",
f"{n_sub}{limit_str}",
sec_str(t_win),
f"[dim]{sec_str(t_io)}[/]",
sec_str(t_jax),
f"[dim]{sec_str(t_py)}[/]",
)
)
self._progress.advance(self._task_id)
n_done = win + 1
total = self._n_windows if self._n_windows is not None else n_done
n_left = total - n_done
self._progress.update(
self._task_id,
description=f"[cyan]{n_done} done, {n_left} left[/]",
)
self._live.update(Group(self._make_table(), self._progress))
def _rich_summary(self, t_total: float):
if self._live is None or t_total <= 0:
return
from rich.console import Group
pct = lambda t: f"{100 * t / t_total:.0f}%" # noqa: E731
cumul = self._cumul_stopped
reason_str = _reason_str(self._cumul_reason_tally)
if cumul:
stop_cell = f"[red]{cumul}[/]\n[dim]{reason_str}[/]"
else:
stop_cell = ""
self._summary = (
"∑",
"",
stop_cell,
"",
f"[bold]{sec_str(t_total)}[/]",
f"[dim]{sec_str(self._t_io)} {pct(self._t_io)}[/]",
f"[bold]{sec_str(self._t_jax)}[/] {pct(self._t_jax)}",
f"[dim]{sec_str(self._t_py)} {pct(self._t_py)}[/]",
)
self._progress.update(
self._task_id,
description=f"[green]{self._wins_done} done, 0 left[/]",
)
self._live.update(Group(self._make_table(), self._progress))
