From c2e5f1010b9fd26474cf10ab0b215c0fb2a5865b Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 16:47:50 +0300 Subject: [PATCH 01/11] feat(memory2): per-stream payload sizes in summary; move summary CLI to dimos mem MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BlobStore grows an optional size_bytes(stream) capability (None = not cheaply knowable): sqlite answers via SUM(LENGTH(data)) — record-header only, no payload page reads (0.2s on a 9.5 GB db) — and the file store via stat. Backend delegates to its blob store, and Stream.summary() appends the size for unfiltered stream queries. The summary verb moves from `dimos map` to `dimos mem` (its module now lives in memory2/cli with imports deferred so `dimos --help` stays fast). --- dimos/mapping/utils/cli/dataset_validation.md | 2 +- dimos/mapping/utils/cli/replay.py | 2 +- dimos/mapping/utils/cli/test_cli.py | 10 +++++--- dimos/memory2/backend.py | 4 +++ dimos/memory2/blobstore/base.py | 4 +++ dimos/memory2/blobstore/file.py | 7 ++++++ dimos/memory2/blobstore/sqlite.py | 11 ++++++++ dimos/memory2/blobstore/test_blobstore.py | 9 +++++++ dimos/memory2/cli/app.py | 3 +++ .../{mapping/utils => memory2}/cli/summary.py | 25 +++++++++++++------ dimos/memory2/stream.py | 18 +++++++++++-- dimos/robot/cli/dimos.py | 2 -- 12 files changed, 79 insertions(+), 18 deletions(-) rename dimos/{mapping/utils => memory2}/cli/summary.py (70%) diff --git a/dimos/mapping/utils/cli/dataset_validation.md b/dimos/mapping/utils/cli/dataset_validation.md index 5f32853ddf..9e4873ed91 100644 --- a/dimos/mapping/utils/cli/dataset_validation.md +++ b/dimos/mapping/utils/cli/dataset_validation.md @@ -1,7 +1,7 @@ Dataset Validation ```sh -dimos map summary recording_go2_mid360_2026-05-29_4-45pm-PST.db +dimos mem summary recording_go2_mid360_2026-05-29_4-45pm-PST.db Stream("color_image"): 11141 items, 2026-05-29 23:32:57 — 2026-05-29 23:45:57 (780.1s) Stream("fastlio_lidar"): 7240 items, 2026-05-29 23:32:56 — 2026-05-29 23:45:57 (781.7s) diff --git a/dimos/mapping/utils/cli/replay.py b/dimos/mapping/utils/cli/replay.py index c78704f725..a69ba28507 100644 --- a/dimos/mapping/utils/cli/replay.py +++ b/dimos/mapping/utils/cli/replay.py @@ -206,8 +206,8 @@ def main( ), ) -> None: """Dump a recording to .rrd (lidar clouds + camera frames) and open it in rerun.""" - from dimos.mapping.utils.cli.summary import _stream_payload_types from dimos.mapping.voxels import VoxelMapTransformer + from dimos.memory2.cli.summary import _stream_payload_types from dimos.memory2.store.sqlite import SqliteStore from dimos.memory2.transform import throttle from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped diff --git a/dimos/mapping/utils/cli/test_cli.py b/dimos/mapping/utils/cli/test_cli.py index cb9ece1f3c..208d7b71b7 100644 --- a/dimos/mapping/utils/cli/test_cli.py +++ b/dimos/mapping/utils/cli/test_cli.py @@ -66,8 +66,8 @@ def _turbojpeg_available() -> bool: ) -def _run(*args: str, timeout: float = 300.0) -> SimpleNamespace: - """Invoke `dimos map ` in-process and capture its result. +def _run(*args: str, timeout: float = 300.0, group: str = "map") -> SimpleNamespace: + """Invoke `dimos ` in-process and capture its result. Uses Typer's CliRunner so the dimos import cost is paid once (module import) rather than per case. `timeout` is kept for call-site compatibility but is a @@ -76,7 +76,7 @@ def _run(*args: str, timeout: float = 300.0) -> SimpleNamespace: """ from dimos.robot.cli.dimos import main as cli_app - res = _runner.invoke(cli_app, ["map", *args]) + res = _runner.invoke(cli_app, [group, *args]) err = res.output if res.exception is not None and not isinstance(res.exception, SystemExit): err += "\n" + "".join(traceback.format_exception(res.exception)) @@ -103,10 +103,12 @@ def dataset() -> str: def test_summary(dataset: str) -> None: - res = _run("summary", dataset) + # summary lives under `dimos mem`, but drives the same recorded dataset. + res = _run("summary", dataset, group="mem") assert res.returncode == 0, res.stderr assert "lidar" in res.stdout assert "odom" in res.stdout + assert "iB" in res.stdout # payload sizes (KiB/MiB/GiB) are included @requires_turbojpeg diff --git a/dimos/memory2/backend.py b/dimos/memory2/backend.py index 7b95bd6335..e96075d929 100644 --- a/dimos/memory2/backend.py +++ b/dimos/memory2/backend.py @@ -76,6 +76,10 @@ def start(self) -> None: def name(self) -> str: return self.metadata_store.name + def size_bytes(self) -> int | None: + """Total stored payload bytes for this stream, or None if not cheaply knowable.""" + return self.blob_store.size_bytes(self.name) if self.blob_store is not None else None + def _make_loader(self, row_id: int) -> Any: bs = self.blob_store if bs is None: diff --git a/dimos/memory2/blobstore/base.py b/dimos/memory2/blobstore/base.py index 36601c5f10..e4fbc1a272 100644 --- a/dimos/memory2/blobstore/base.py +++ b/dimos/memory2/blobstore/base.py @@ -54,5 +54,9 @@ def delete(self, stream_name: str, key: int) -> None: """Delete a blob by stream name and observation id.""" ... + def size_bytes(self, stream_name: str) -> int | None: + """Total stored payload bytes for a stream, or None if not cheaply knowable.""" + return None + def serialize(self) -> dict[str, Any]: return {"class": qual(type(self)), "config": self.config.model_dump()} diff --git a/dimos/memory2/blobstore/file.py b/dimos/memory2/blobstore/file.py index 452b13c9b7..edfcf26f19 100644 --- a/dimos/memory2/blobstore/file.py +++ b/dimos/memory2/blobstore/file.py @@ -67,3 +67,10 @@ def delete(self, stream_name: str, key: int) -> None: p.unlink() except FileNotFoundError: raise KeyError(f"No blob for stream={stream_name!r}, key={key}") from None + + def size_bytes(self, stream_name: str) -> int | None: + validate_identifier(stream_name) + stream_dir = self._root / stream_name + if not stream_dir.is_dir(): + return 0 + return sum(p.stat().st_size for p in stream_dir.iterdir()) diff --git a/dimos/memory2/blobstore/sqlite.py b/dimos/memory2/blobstore/sqlite.py index 85b539bfd6..06021dcd9a 100644 --- a/dimos/memory2/blobstore/sqlite.py +++ b/dimos/memory2/blobstore/sqlite.py @@ -105,3 +105,14 @@ def delete(self, stream_name: str, key: int) -> None: raise KeyError(f"No blob for stream={stream_name!r}, key={key}") from None if cur.rowcount == 0: raise KeyError(f"No blob for stream={stream_name!r}, key={key}") + + def size_bytes(self, stream_name: str) -> int | None: + # LENGTH(data) comes from the record header, so this never reads payload pages. + validate_identifier(stream_name) + try: + row = self._conn.execute( + f'SELECT SUM(LENGTH(data)) FROM "{stream_name}_blob"' + ).fetchone() + except sqlite3.OperationalError: # no blob table for this stream + return 0 + return int(row[0]) if row[0] is not None else 0 diff --git a/dimos/memory2/blobstore/test_blobstore.py b/dimos/memory2/blobstore/test_blobstore.py index ade6aa4cc6..2cae1a8a55 100644 --- a/dimos/memory2/blobstore/test_blobstore.py +++ b/dimos/memory2/blobstore/test_blobstore.py @@ -60,3 +60,12 @@ def test_large_blob(self, blob_store: BlobStore) -> None: blob_store.put("big", 0, data) assert blob_store.get("big", 0) == data assert blob_store.get("big", 0) == data + + def test_size_bytes(self, blob_store: BlobStore) -> None: + assert blob_store.size_bytes("s") == 0 + blob_store.put("s", 1, b"12345") + blob_store.put("s", 2, b"1234567890") + assert blob_store.size_bytes("s") == 15 + blob_store.delete("s", 1) + assert blob_store.size_bytes("s") == 10 + assert blob_store.size_bytes("other") == 0 diff --git a/dimos/memory2/cli/app.py b/dimos/memory2/cli/app.py index 2424c600df..3fbdb38282 100644 --- a/dimos/memory2/cli/app.py +++ b/dimos/memory2/cli/app.py @@ -18,7 +18,10 @@ import typer +from dimos.memory2.cli.summary import main as _summary_main + mem_app = typer.Typer(help="memory2 store commands", no_args_is_help=True) +mem_app.command("summary")(_summary_main) @mem_app.command() diff --git a/dimos/mapping/utils/cli/summary.py b/dimos/memory2/cli/summary.py similarity index 70% rename from dimos/mapping/utils/cli/summary.py rename to dimos/memory2/cli/summary.py index 694565b6e8..cb335c9b28 100644 --- a/dimos/mapping/utils/cli/summary.py +++ b/dimos/memory2/cli/summary.py @@ -15,24 +15,30 @@ """Print ``Store.summary()`` for a memory2 sqlite recording. Usage: - uv run dimos map summary mid360 + uv run dimos mem summary mid360 """ from __future__ import annotations -import json -from pathlib import Path -import sqlite3 +from typing import TYPE_CHECKING import typer -from dimos.memory2.codecs.base import _resolve_payload_type -from dimos.memory2.store.sqlite import SqliteStore -from dimos.utils.data import resolve_named_path +if TYPE_CHECKING: + from pathlib import Path + +# Heavy dimos imports (memory2 store → codecs, msgs) are deferred into the +# function bodies so that `dimos --help` — which imports this module just to +# register the `mem summary` command — stays fast. See test_cli_startup.py. def _stream_payload_types(db_path: Path) -> dict[str, type]: """Read each stream's registered payload type from the _streams table.""" + import json + import sqlite3 + + from dimos.memory2.codecs.base import _resolve_payload_type + conn = sqlite3.connect(str(db_path)) try: rows = conn.execute("SELECT name, config FROM _streams").fetchall() @@ -44,7 +50,10 @@ def _stream_payload_types(db_path: Path) -> dict[str, type]: def main( dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), ) -> None: - """Print per-stream counts and time ranges for a recorded SQLite dataset.""" + """Print per-stream counts, time ranges, and payload sizes for a recorded SQLite dataset.""" + from dimos.memory2.store.sqlite import SqliteStore + from dimos.utils.data import resolve_named_path + db_path = resolve_named_path(dataset, ".db") payload_types = _stream_payload_types(db_path) diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index c3cb1a17ae..af7164d329 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -497,9 +497,23 @@ def get_time_range(self) -> tuple[float, float]: return (first.ts, last.ts) def summary(self) -> str: - """Return a short human-readable summary: count, time range, avg frequency.""" + """Return a short human-readable summary: count, time range, avg frequency, size.""" from datetime import datetime, timezone + from dimos.memory2.backend import Backend + from dimos.utils.human import human_bytes + + # Size is whole-stream (blob totals), so only report it when the query + # doesn't narrow the stream. + unfiltered = not self._query.filters and self._query.limit_val is None + size = "" + if ( + unfiltered + and isinstance(self._source, Backend) + and (sz := self._source.size_bytes()) is not None + ): + size = f", {human_bytes(sz)}" + n = self.count() if n == 0: return f"{self}: empty" @@ -510,7 +524,7 @@ def summary(self) -> str: dt1 = datetime.fromtimestamp(t1, tz=timezone.utc).strftime(fmt) dur = t1 - t0 hz = f", {(n - 1) / dur:.2f} Hz" if dur > 0 else "" - return f"{self}: {n} items, {dt0} — {dt1} ({dur:.1f}s{hz})" + return f"{self}: {n} items, {dt0} — {dt1} ({dur:.1f}s{hz}{size})" def materialize(self) -> Stream[T, O]: """Materialize into memory and return a replayable stream. diff --git a/dimos/robot/cli/dimos.py b/dimos/robot/cli/dimos.py index d7615a2795..6c7cc18226 100644 --- a/dimos/robot/cli/dimos.py +++ b/dimos/robot/cli/dimos.py @@ -42,7 +42,6 @@ from dimos.mapping.utils.cli.rename import main as _map_rename_main from dimos.mapping.utils.cli.replay import main as _map_replay_main from dimos.mapping.utils.cli.replay_marker import main as _map_replay_marker_main -from dimos.mapping.utils.cli.summary import main as _map_summary_main from dimos.robot.unitree.go2.cli.go2tool import app as go2tool_app from dimos.utils.logging_config import setup_logger from dimos.visualization.rerun.constants import RerunOpenOption @@ -722,7 +721,6 @@ def dataprep_inspect( inspect(dataset, cast("Literal['lerobot', 'hdf5'] | None", output_format)) -map_app.command("summary")(_map_summary_main) map_app.command("rename")(_map_rename_main) map_app.command("pose-fill")(_map_pose_fill_main) map_app.command("replay")(_map_replay_main) From e3e0603f251879d2779f3a339561e490a0cfa0df Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 16:55:48 +0300 Subject: [PATCH 02/11] feat(memory2): mem summary renders a rich table sorted by size One row per stream plus a totals row; Items/Hz/Size are heat-colored with the pubsub-benchmark red->green ANSI gradient, log-scaled per column. Stream.size_bytes() becomes public API (None when the query narrows the stream), and piped output gets a wide console so long stream names don't truncate. --- dimos/memory2/cli/summary.py | 78 ++++++++++++++++++++++++++++++++++-- dimos/memory2/stream.py | 23 ++++++----- 2 files changed, 87 insertions(+), 14 deletions(-) diff --git a/dimos/memory2/cli/summary.py b/dimos/memory2/cli/summary.py index cb335c9b28..c7198d8d7b 100644 --- a/dimos/memory2/cli/summary.py +++ b/dimos/memory2/cli/summary.py @@ -20,13 +20,15 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import typer if TYPE_CHECKING: from pathlib import Path + from dimos.memory2.stream import Stream + # Heavy dimos imports (memory2 store → codecs, msgs) are deferred into the # function bodies so that `dimos --help` — which imports this module just to # register the `mem summary` command — stays fast. See test_cli_startup.py. @@ -47,21 +49,91 @@ def _stream_payload_types(db_path: Path) -> dict[str, type]: return {name: _resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} +# ANSI 256 gradient: red -> orange -> yellow -> green (same as the pubsub +# benchmark heatmaps in dimos/protocol/pubsub/benchmark/type.py). +_GRADIENT = [52, 88, 124, 160, 196, 202, 208, 214, 220, 226, 190, 154, 148, 118, 82, 46, 40, 34] + + +def _shade(value: float, lo: float, hi: float) -> str: + """Rich style for ``value`` relative to [lo, hi], log-scaled (columns span decades).""" + from math import log + + if value <= 0: + return "dim" + t = 0.5 if hi <= lo else (log(value) - log(lo)) / (log(hi) - log(lo)) + return f"color({_GRADIENT[round(t * (len(_GRADIENT) - 1))]})" + + +def _heat(text: str, value: float, column: list[float]) -> str: + """Wrap ``text`` in rich markup colored by ``value``'s rank within ``column``.""" + positive = [v for v in column if v > 0] + lo, hi = (min(positive), max(positive)) if positive else (0.0, 0.0) + return f"[{_shade(value, lo, hi)}]{text}[/]" + + def main( dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), ) -> None: """Print per-stream counts, time ranges, and payload sizes for a recorded SQLite dataset.""" + from datetime import datetime, timezone + + from rich.console import Console + from rich.table import Table + from dimos.memory2.store.sqlite import SqliteStore from dimos.utils.data import resolve_named_path + from dimos.utils.human import human_bytes db_path = resolve_named_path(dataset, ".db") payload_types = _stream_payload_types(db_path) + rows: list[tuple[str, int, float | None, float | None, int]] = [] store = SqliteStore(path=str(db_path)) with store: for name, ptype in payload_types.items(): - store.stream(name, ptype) - print(store.summary()) + stream: Stream[Any] = store.stream(name, ptype) + n = stream.count() + t0, t1 = stream.get_time_range() if n else (None, None) + rows.append((name, n, t0, t1, stream.size_bytes() or 0)) + rows.sort(key=lambda r: r[4], reverse=True) + + table = Table(title=db_path.name) + table.add_column("Stream", style="cyan") + table.add_column("Items", justify="right") + table.add_column("Hz", justify="right") + table.add_column("Start (UTC)") + table.add_column("Duration", justify="right") + table.add_column("Size", justify="right") + + def hz(n: int, t0: float | None, t1: float | None) -> float: + return (n - 1) / (t1 - t0) if t0 is not None and t1 is not None and t1 > t0 else 0.0 + + items_col = [float(r[1]) for r in rows] + hz_col = [hz(r[1], r[2], r[3]) for r in rows] + size_col = [float(r[4]) for r in rows] + + for name, n, t0, t1, size in rows: + dur = t1 - t0 if t0 is not None and t1 is not None else None + rate = hz(n, t0, t1) + table.add_row( + name, + _heat(f"{n:,}", n, items_col), + _heat(f"{rate:.1f}", rate, hz_col) if rate > 0 else "—", + datetime.fromtimestamp(t0, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S") + if t0 is not None + else "—", + f"{dur:.1f}s" if dur is not None else "—", + _heat(human_bytes(size), size, size_col), + ) + table.add_section() + table.add_row( + "total", f"{sum(r[1] for r in rows):,}", "", "", "", human_bytes(sum(r[4] for r in rows)) + ) + + console = Console() + if not console.is_terminal: # piped: don't squeeze the table into the 80-col default + console = Console(width=250) + console.print(table) if __name__ == "__main__": diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index af7164d329..278c022145 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -496,23 +496,24 @@ def get_time_range(self) -> tuple[float, float]: last = self.last() return (first.ts, last.ts) + def size_bytes(self) -> int | None: + """Total stored payload bytes, or None if unknown or the query narrows the stream.""" + from dimos.memory2.backend import Backend + + if self._query.filters or self._query.limit_val is not None: + return None + if isinstance(self._source, Backend): + return self._source.size_bytes() + return None + def summary(self) -> str: """Return a short human-readable summary: count, time range, avg frequency, size.""" from datetime import datetime, timezone - from dimos.memory2.backend import Backend from dimos.utils.human import human_bytes - # Size is whole-stream (blob totals), so only report it when the query - # doesn't narrow the stream. - unfiltered = not self._query.filters and self._query.limit_val is None - size = "" - if ( - unfiltered - and isinstance(self._source, Backend) - and (sz := self._source.size_bytes()) is not None - ): - size = f", {human_bytes(sz)}" + sz = self.size_bytes() + size = f", {human_bytes(sz)}" if sz is not None else "" n = self.count() if n == 0: From 8f6dbaa972f4427a629905d8c50b31c1166f6547 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 17:01:11 +0300 Subject: [PATCH 03/11] fix(go2): replay resolves go2_lidar/go2_odom with lidar/odom fallback mid360 recordings (e.g. china_office) name the raw robot streams go2_lidar/go2_odom; older datasets (go2_hongkong_office) use lidar/odom. ReplayConnection now picks whichever the dataset has and reports the available streams when neither exists. --- dimos/robot/unitree/go2/connection.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/dimos/robot/unitree/go2/connection.py b/dimos/robot/unitree/go2/connection.py index b2f74bf327..fe3ee2a347 100644 --- a/dimos/robot/unitree/go2/connection.py +++ b/dimos/robot/unitree/go2/connection.py @@ -32,7 +32,7 @@ from dimos.core.module import Module, ModuleConfig from dimos.core.resource import CompositeResource from dimos.core.stream import In, Out -from dimos.memory2.replay import Replay, resolve_db_path +from dimos.memory2.replay import Replay, ReplayStream, resolve_db_path from dimos.memory2.store.sqlite import SqliteStore from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped from dimos.msgs.geometry_msgs.Quaternion import Quaternion @@ -188,13 +188,29 @@ def set_motion_mode(self, name: str) -> None: def set_rage_mode(self, enable: bool) -> bool: return True + def _stream_name(self, *names: str) -> str: + """Return the first of ``names`` present in the dataset (stream naming + changed over time: mid360 recordings use go2_lidar/go2_odom, older ones + lidar/odom).""" + available = self.replay.list_streams() + for name in names: + if name in available: + return name + raise KeyError(f"None of {names!r} in dataset {self.dataset!r}; available: {available}") + @simple_mcache def lidar_stream(self) -> Observable[PointCloud2]: - return self.replay.streams.lidar.observable() + stream: ReplayStream[PointCloud2] = self.replay.stream( + self._stream_name("go2_lidar", "lidar") + ) + return stream.observable() @simple_mcache def odom_stream(self) -> Observable[PoseStamped]: - return self.replay.streams.odom.observable() + stream: ReplayStream[PoseStamped] = self.replay.stream( + self._stream_name("go2_odom", "odom") + ) + return stream.observable() @simple_mcache def video_stream(self) -> Observable[Image]: From a3c3619f5da8cedf690c404f5cfea9eaa43adae6 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 17:05:26 +0300 Subject: [PATCH 04/11] feat(memory2): mem summary progress bar; mem rerun resolves names like other db verbs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Scanning big datasets no longer looks frozen — a transient rich progress bar names the stream being scanned. open_store() goes through resolve_named_path (cwd -> data/ -> LFS pull) instead of requiring an on-disk path, defaulting to .db unless the name says .mcap. --- dimos/memory2/cli/app.py | 2 +- dimos/memory2/cli/render.py | 19 +++++++++++++------ dimos/memory2/cli/summary.py | 6 +++++- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/dimos/memory2/cli/app.py b/dimos/memory2/cli/app.py index 3fbdb38282..8c4278a207 100644 --- a/dimos/memory2/cli/app.py +++ b/dimos/memory2/cli/app.py @@ -26,7 +26,7 @@ @mem_app.command() def rerun( - path: str = typer.Argument(..., help="Store to render (.mcap path/name or .db)"), + path: str = typer.Argument(..., help="Store: bare name (cwd, data/, LFS), .db or .mcap path"), out: str = typer.Option(None, "--out", help="Output .rrd (default: alongside the source)"), seconds: float = typer.Option(None, "--seconds", help="Only the first N seconds"), no_gui: bool = typer.Option(False, "--no-gui", help="Write the .rrd but don't open the viewer"), diff --git a/dimos/memory2/cli/render.py b/dimos/memory2/cli/render.py index 0869c7c45f..e34b60ebfd 100644 --- a/dimos/memory2/cli/render.py +++ b/dimos/memory2/cli/render.py @@ -33,14 +33,21 @@ def open_store(path: str) -> Store: - """Open a store by file type (``.db`` -> SqliteStore, else Go2 mcap).""" - if str(path).endswith(".db"): - from dimos.memory2.store.sqlite import SqliteStore + """Open a store by name or path (``.db`` -> SqliteStore, ``.mcap`` -> Go2 mcap). - return SqliteStore(path=path, must_exist=True) - from dimos.robot.unitree.go2.dds.store import Go2McapStore # lazy: robot-layer codec set + Bare names resolve like the other db verbs: cwd, then ``data/``, then an + LFS pull — defaulting to ``.db`` unless the name says ``.mcap``. + """ + from dimos.utils.data import resolve_named_path + + resolved = resolve_named_path(path, ".mcap" if str(path).endswith(".mcap") else ".db") + if resolved.suffix == ".mcap": + from dimos.robot.unitree.go2.dds.store import Go2McapStore # lazy: robot-layer codec set + + return Go2McapStore(path=str(resolved)) + from dimos.memory2.store.sqlite import SqliteStore - return Go2McapStore(path=path) + return SqliteStore(path=str(resolved), must_exist=True) def _open_viewer(rrd: str) -> None: diff --git a/dimos/memory2/cli/summary.py b/dimos/memory2/cli/summary.py index c7198d8d7b..95b111ce3f 100644 --- a/dimos/memory2/cli/summary.py +++ b/dimos/memory2/cli/summary.py @@ -78,6 +78,7 @@ def main( from datetime import datetime, timezone from rich.console import Console + from rich.progress import Progress from rich.table import Table from dimos.memory2.store.sqlite import SqliteStore @@ -89,12 +90,15 @@ def main( rows: list[tuple[str, int, float | None, float | None, int]] = [] store = SqliteStore(path=str(db_path)) - with store: + with store, Progress(transient=True) as prog: + task = prog.add_task("scanning", total=len(payload_types)) for name, ptype in payload_types.items(): + prog.update(task, description=name) stream: Stream[Any] = store.stream(name, ptype) n = stream.count() t0, t1 = stream.get_time_range() if n else (None, None) rows.append((name, n, t0, t1, stream.size_bytes() or 0)) + prog.advance(task) rows.sort(key=lambda r: r[4], reverse=True) table = Table(title=db_path.name) From 92f75c3469cc9363446d54474fd9ad8d42e4e714 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 17:19:04 +0300 Subject: [PATCH 05/11] refactor(memory2): shared per-observation progress renders a rich bar progress() keeps its Stream.tap callback shape but now returns a ProgressBar backed by ONE process-wide rich Progress, so sequential or interleaved streams compose instead of fighting over the terminal. Completed bars persist as a plain line; close() finalizes a bar abandoned early (mem rerun's --seconds window). map replay drops its private copy of the old print-based progress. --- dimos/mapping/utils/cli/replay.py | 31 +------- dimos/memory2/cli/render.py | 2 +- dimos/memory2/utils/progress.py | 122 +++++++++++++++++++++++------- 3 files changed, 97 insertions(+), 58 deletions(-) diff --git a/dimos/mapping/utils/cli/replay.py b/dimos/mapping/utils/cli/replay.py index a69ba28507..c9d3995041 100644 --- a/dimos/mapping/utils/cli/replay.py +++ b/dimos/mapping/utils/cli/replay.py @@ -34,7 +34,6 @@ from collections.abc import Callable from pathlib import Path import subprocess -import time from typing import TYPE_CHECKING, Any import rerun as rr @@ -52,34 +51,10 @@ def _progress(total: int, label: str) -> Callable[[Observation[Any]], None]: - """Matches dimos/mapping/utils/cli/map.py:progress.""" - seen = 0 - wall_start: float | None = None - last_wall: float | None = None - first_ts: float | None = None - - def tick(obs: Observation[Any]) -> None: - nonlocal seen, wall_start, last_wall, first_ts - now = time.monotonic() - if wall_start is None: - wall_start = now - first_ts = obs.ts - assert first_ts is not None - frame_ms = (now - last_wall) * 1000 if last_wall is not None else 0.0 - last_wall = now - seen += 1 - pct = 100 * seen // total if total else 100 - wall = now - wall_start - data = obs.ts - first_ts - speed = data / wall if wall > 0 else 0.0 - end = "\n" if seen >= total else "" - print( - f"\r{label} {pct:>3}% [{seen}/{total}] {data:.1f}s ({speed:.1f} x rt) {frame_ms:.0f}ms/frame", - end=end, - flush=True, - ) + """Deferred alias for the shared rich progress bar (keeps module import light).""" + from dimos.memory2.utils.progress import progress - return tick + return progress(total, label) def _log_clouds( diff --git a/dimos/memory2/cli/render.py b/dimos/memory2/cli/render.py index e34b60ebfd..615fa77546 100644 --- a/dimos/memory2/cli/render.py +++ b/dimos/memory2/cli/render.py @@ -112,7 +112,7 @@ def render_store( report = progress(stream.count(), label=name) for obs in stream: if seconds is not None and obs.ts - t0 > seconds: - print() # terminate the windowed (sub-100%) progress line + report.close() # finalize the windowed (sub-100%) bar break if obs.data is None: # e.g. a truncated/corrupt frame that failed to decode report(obs) diff --git a/dimos/memory2/utils/progress.py b/dimos/memory2/utils/progress.py index 29aaf8d4e4..c0e816b17b 100644 --- a/dimos/memory2/utils/progress.py +++ b/dimos/memory2/utils/progress.py @@ -12,51 +12,115 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""A per-observation progress callback, shaped for ``Stream.tap``. +"""A per-observation progress bar, shaped for ``Stream.tap``. stream.tap(progress(stream.count(), "render")).drain() -Prints a single rewritten line: percent, count, data-seconds covered, speed -relative to wall-clock (``x rt``), and per-frame latency. +Renders a rich progress bar: label, percent, count, data-seconds covered, +speed relative to wall-clock (``x rt``), and per-frame latency. On completion +the bar persists as one plain line. All bars share a single live display, so +sequential or interleaved streams compose; call :meth:`ProgressBar.close` when +abandoning a stream early (e.g. a bounded time window). """ from __future__ import annotations -from collections.abc import Callable +import atexit import time from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + from rich.progress import Progress, TaskID + from dimos.memory2.type.observation import Observation +_REFRESH_S = 0.1 # manual refresh cap; keeps 200 Hz streams from redrawing per frame + +_shared: Progress | None = None + -def progress(total: int, label: str = "") -> Callable[[Observation[Any]], None]: - """Return a callback that prints streaming progress, one call per observation.""" - seen = 0 - wall_start: float | None = None - last_wall: float | None = None - first_ts: float | None = None +def _shared_progress() -> Progress: + """One Progress (one live display) shared by every bar in the process.""" + global _shared + if _shared is None: + from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + TaskProgressColumn, + TextColumn, + ) + + # auto_refresh=False: no background render thread to leak if a stream + # ends short of ``total`` — bars refresh (throttled) from the callback. + # transient: close() prints its own persisted line, the live bar clears. + _shared = Progress( + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TaskProgressColumn(), + MofNCompleteColumn(), + TextColumn("{task.fields[stats]}", style="dim"), + auto_refresh=False, + transient=True, + ) + atexit.register(_shared.stop) # restore the cursor if a bar is abandoned + return _shared - def _progress(obs: Observation[Any]) -> None: - nonlocal seen, wall_start, last_wall, first_ts + +class ProgressBar: + """Callable per-observation progress; ``close()`` finalizes an unfinished bar.""" + + def __init__(self, total: int, label: str = "") -> None: + self._total = total + self._label = label or "processing" + self._prog = _shared_progress() + self._task: TaskID | None = self._prog.add_task(self._label, total=total, stats="") + self._seen = 0 + self._wall_start: float | None = None + self._last_wall: float | None = None + self._last_refresh: float | None = None + self._first_ts: float | None = None + self._stats = "" + + def __call__(self, obs: Observation[Any]) -> None: + if self._task is None: + return # closed; observations may keep flowing (e.g. a tap after close) now = time.monotonic() - if wall_start is None: - wall_start = now - first_ts = obs.ts - assert first_ts is not None # narrowed by the same `if` above - frame_ms = (now - last_wall) * 1000 if last_wall is not None else 0.0 - last_wall = now - seen += 1 - pct = 100 * seen // total if total else 100 - wall = now - wall_start - data = obs.ts - first_ts + if self._wall_start is None: + self._wall_start = now + self._first_ts = obs.ts + self._prog.start() + assert self._first_ts is not None # narrowed by the same `if` above + frame_ms = (now - self._last_wall) * 1000 if self._last_wall is not None else 0.0 + self._last_wall = now + self._seen += 1 + wall = now - self._wall_start + data = obs.ts - self._first_ts speed = data / wall if wall > 0 else 0.0 - end = "\n" if seen >= total else "" - prefix = f"{label} " if label else "" - print( - f"\r{prefix}{pct:>3}% [{seen}/{total}] {data:.1f}s ({speed:.1f} x rt) {frame_ms:.0f}ms/frame", - end=end, - flush=True, + self._stats = f"{data:.1f}s ({speed:.1f} x rt) {frame_ms:.0f}ms/frame" + self._prog.update(self._task, advance=1, stats=self._stats) + if self._seen >= self._total: + self.close() + elif self._last_refresh is None or now - self._last_refresh >= _REFRESH_S: + self._last_refresh = now + self._prog.refresh() + + def close(self) -> None: + """Persist the bar as a plain line and release it from the live display.""" + if self._task is None: + return # already closed + self._prog.remove_task(self._task) # before print: Live re-renders on print + self._task = None + pct = 100 * self._seen // self._total if self._total else 100 + self._prog.console.print( + f"{self._label} {pct}% [{self._seen}/{self._total}] {self._stats}", + markup=False, + highlight=False, ) + if not self._prog.tasks: + self._prog.stop() + - return _progress +def progress(total: int, label: str = "") -> ProgressBar: + """Return a callback that renders streaming progress, one call per observation.""" + return ProgressBar(total, label) From 693fdf8c031943d4fcc2fa2a168220769257c46e Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 18:15:29 +0300 Subject: [PATCH 06/11] refactor(mapping): replay imports the shared progress directly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The deferral wrapper bought nothing — memory2/utils/progress.py is a light module (rich is deferred inside it), so import it at the top. --- dimos/mapping/utils/cli/replay.py | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/dimos/mapping/utils/cli/replay.py b/dimos/mapping/utils/cli/replay.py index c9d3995041..b056f44321 100644 --- a/dimos/mapping/utils/cli/replay.py +++ b/dimos/mapping/utils/cli/replay.py @@ -31,7 +31,6 @@ from __future__ import annotations -from collections.abc import Callable from pathlib import Path import subprocess from typing import TYPE_CHECKING, Any @@ -39,24 +38,18 @@ import rerun as rr import typer +from dimos.memory2.utils.progress import progress + # Heavy dimos imports (mapping/memory2 → torch, scipy, open3d) are deferred into # main() so that `dimos map --help` stays fast. See test_cli_startup.py and the # same pattern in dimos/mapping/utils/cli/map.py. if TYPE_CHECKING: from dimos.memory2.stream import Stream - from dimos.memory2.type.observation import Observation from dimos.msgs.sensor_msgs.PointCloud2 import PointCloud2 TIMELINE = "ts" -def _progress(total: int, label: str) -> Callable[[Observation[Any]], None]: - """Deferred alias for the shared rich progress bar (keeps module import light).""" - from dimos.memory2.utils.progress import progress - - return progress(total, label) - - def _log_clouds( label: str, stream: Stream[PointCloud2], @@ -74,7 +67,7 @@ def _log_clouds( whole pipeline. """ n = total if total is not None else stream.count() - cb = _progress(n, label) + cb = progress(n, label) for obs in stream: cb(obs) rr.set_time(TIMELINE, timestamp=obs.ts) @@ -97,7 +90,7 @@ def _log_path( without a pose are skipped. """ n = stream.count() - cb = _progress(n, label) + cb = progress(n, label) points: list[tuple[float, float, float]] = [] last_ts: float | None = None emit_count = 0 @@ -284,7 +277,7 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: # fastlio pose axis + path from fastlio_odometry stream. if "fastlio_odometry" in store.streams: odometry = clipped("fastlio_odometry", Odometry) - cb = _progress(odometry.count(), "fastlio_odometry") + cb = progress(odometry.count(), "fastlio_odometry") for obs in odometry: cb(obs) if obs.pose_tuple is None: @@ -308,7 +301,7 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: # Go2 native odom pose axis + path. if "odom" in store.streams: odom = clipped("odom", PoseStamped) - cb = _progress(odom.count(), " odom") + cb = progress(odom.count(), " odom") for odom_obs in odom: cb(odom_obs) if odom_obs.pose_tuple is None: @@ -334,7 +327,7 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: color_image.transform(throttle(1.0 / camera_hz)) if camera_hz > 0 else color_image ) n_img = cam_pipeline.count() - cb = _progress(n_img, " color_image") + cb = progress(n_img, " color_image") for img_obs in cam_pipeline: cb(img_obs) rr.set_time(TIMELINE, timestamp=img_obs.ts) From 07bcc962b7441b17df53cad806e97b61a148ba95 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 18:18:25 +0300 Subject: [PATCH 07/11] fix(memory2): size_bytes guards all subset queries; file store sums only blobs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review (PR #2730): offset/search queries narrowed the stream but still reported full-stream size — treat them like filters/limit (ordering keeps membership, so it stays allowed). FileBlobStore now sums only {key}.bin files instead of everything in the stream dir. --- dimos/memory2/blobstore/file.py | 2 +- dimos/memory2/stream.py | 10 +++++++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/dimos/memory2/blobstore/file.py b/dimos/memory2/blobstore/file.py index edfcf26f19..137c39b043 100644 --- a/dimos/memory2/blobstore/file.py +++ b/dimos/memory2/blobstore/file.py @@ -73,4 +73,4 @@ def size_bytes(self, stream_name: str) -> int | None: stream_dir = self._root / stream_name if not stream_dir.is_dir(): return 0 - return sum(p.stat().st_size for p in stream_dir.iterdir()) + return sum(p.stat().st_size for p in stream_dir.glob("*.bin") if p.is_file()) diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index 278c022145..976e47afef 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -500,7 +500,15 @@ def size_bytes(self) -> int | None: """Total stored payload bytes, or None if unknown or the query narrows the stream.""" from dimos.memory2.backend import Backend - if self._query.filters or self._query.limit_val is not None: + q = self._query + narrowed = ( + q.filters + or q.limit_val is not None + or q.offset_val + or q.search_vec is not None + or q.search_text is not None + ) + if narrowed: return None if isinstance(self._source, Backend): return self._source.size_bytes() From e90c5bb77cbb7ba274e56e9842caca0bd0132489 Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 18:24:14 +0300 Subject: [PATCH 08/11] =?UTF-8?q?refactor(memory2):=20each=20ProgressBar?= =?UTF-8?q?=20owns=20its=20Progress=20=E2=80=94=20drop=20the=20shared=20si?= =?UTF-8?q?ngleton?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Simpler model: one bar live at a time, per-bar transient Progress, and close() is just stop-the-live-bar + a plain persisted print. Starting a new bar finalizes a dangling one (rich allows a single live display). --- dimos/memory2/utils/progress.py | 81 ++++++++++++++++----------------- 1 file changed, 40 insertions(+), 41 deletions(-) diff --git a/dimos/memory2/utils/progress.py b/dimos/memory2/utils/progress.py index c0e816b17b..cf8cea5b0c 100644 --- a/dimos/memory2/utils/progress.py +++ b/dimos/memory2/utils/progress.py @@ -18,9 +18,10 @@ Renders a rich progress bar: label, percent, count, data-seconds covered, speed relative to wall-clock (``x rt``), and per-frame latency. On completion -the bar persists as one plain line. All bars share a single live display, so -sequential or interleaved streams compose; call :meth:`ProgressBar.close` when -abandoning a stream early (e.g. a bounded time window). +the live bar is replaced by one plain persisted line. One bar is live at a +time — starting a new bar finalizes a dangling one (rich allows a single live +display); call :meth:`ProgressBar.close` when abandoning a stream early +(e.g. a bounded time window). """ from __future__ import annotations @@ -30,19 +31,25 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: - from rich.progress import Progress, TaskID - from dimos.memory2.type.observation import Observation -_REFRESH_S = 0.1 # manual refresh cap; keeps 200 Hz streams from redrawing per frame +_REFRESH_S = 0.1 # refresh cap; keeps 200 Hz streams from redrawing per frame + +_current: ProgressBar | None = None + -_shared: Progress | None = None +def _close_current() -> None: + if _current is not None: + _current.close() -def _shared_progress() -> Progress: - """One Progress (one live display) shared by every bar in the process.""" - global _shared - if _shared is None: +atexit.register(_close_current) # restore the cursor if a bar is abandoned + + +class ProgressBar: + """Callable per-observation progress; ``close()`` finalizes an unfinished bar.""" + + def __init__(self, total: int, label: str = "") -> None: from rich.progress import ( BarColumn, MofNCompleteColumn, @@ -51,30 +58,21 @@ def _shared_progress() -> Progress: TextColumn, ) - # auto_refresh=False: no background render thread to leak if a stream - # ends short of ``total`` — bars refresh (throttled) from the callback. + self._total = total + self._label = label or "processing" # transient: close() prints its own persisted line, the live bar clears. - _shared = Progress( + # auto_refresh=False: no render thread; refreshes (throttled) from the callback. + self._prog = Progress( TextColumn("[progress.description]{task.description}"), BarColumn(), TaskProgressColumn(), MofNCompleteColumn(), TextColumn("{task.fields[stats]}", style="dim"), - auto_refresh=False, transient=True, + auto_refresh=False, ) - atexit.register(_shared.stop) # restore the cursor if a bar is abandoned - return _shared - - -class ProgressBar: - """Callable per-observation progress; ``close()`` finalizes an unfinished bar.""" - - def __init__(self, total: int, label: str = "") -> None: - self._total = total - self._label = label or "processing" - self._prog = _shared_progress() - self._task: TaskID | None = self._prog.add_task(self._label, total=total, stats="") + self._task = self._prog.add_task(self._label, total=total, stats="") + self._closed = False self._seen = 0 self._wall_start: float | None = None self._last_wall: float | None = None @@ -83,10 +81,14 @@ def __init__(self, total: int, label: str = "") -> None: self._stats = "" def __call__(self, obs: Observation[Any]) -> None: - if self._task is None: - return # closed; observations may keep flowing (e.g. a tap after close) + if self._closed: + return # observations may keep flowing after close (e.g. a tap downstream) now = time.monotonic() if self._wall_start is None: + global _current + if _current is not None: + _current.close() # only one live display may be active + _current = self self._wall_start = now self._first_ts = obs.ts self._prog.start() @@ -106,19 +108,16 @@ def __call__(self, obs: Observation[Any]) -> None: self._prog.refresh() def close(self) -> None: - """Persist the bar as a plain line and release it from the live display.""" - if self._task is None: - return # already closed - self._prog.remove_task(self._task) # before print: Live re-renders on print - self._task = None + """Clear the live bar and persist it as a plain line.""" + global _current + if self._closed: + return + self._closed = True + if _current is self: + _current = None + self._prog.stop() # transient: wipes the live bar pct = 100 * self._seen // self._total if self._total else 100 - self._prog.console.print( - f"{self._label} {pct}% [{self._seen}/{self._total}] {self._stats}", - markup=False, - highlight=False, - ) - if not self._prog.tasks: - self._prog.stop() + print(f"{self._label} {pct}% [{self._seen}/{self._total}] {self._stats}") def progress(total: int, label: str = "") -> ProgressBar: From fd806f31b68bd6dc524c8a5b26bace55a28b4cfa Mon Sep 17 00:00:00 2001 From: Ivan Nikolic Date: Mon, 6 Jul 2026 18:31:36 +0300 Subject: [PATCH 09/11] fix(memory2): file-store size counts only integer-keyed blobs Review (PR #2730): match put()'s {key}.bin layout exactly so a stray backup.bin or temp file beside real blobs can't inflate the size. --- dimos/memory2/blobstore/file.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dimos/memory2/blobstore/file.py b/dimos/memory2/blobstore/file.py index 137c39b043..60c8997db0 100644 --- a/dimos/memory2/blobstore/file.py +++ b/dimos/memory2/blobstore/file.py @@ -73,4 +73,7 @@ def size_bytes(self, stream_name: str) -> int | None: stream_dir = self._root / stream_name if not stream_dir.is_dir(): return 0 - return sum(p.stat().st_size for p in stream_dir.glob("*.bin") if p.is_file()) + # Only integer-keyed files — the layout put() writes. + return sum( + p.stat().st_size for p in stream_dir.glob("*.bin") if p.is_file() and p.stem.isdigit() + ) From 9c319d1d92dcfc55f1fd9f384e4c02195e6f23dd Mon Sep 17 00:00:00 2001 From: Paul Nechifor Date: Tue, 7 Jul 2026 01:40:52 +0300 Subject: [PATCH 10/11] Auto-fixes for feat/ivan/mem-summary-size (#2746) --- dimos/mapping/utils/cli/dataset_validation.md | 10 ++---- dimos/mapping/utils/cli/rename.py | 4 +-- dimos/mapping/utils/cli/replay.py | 4 +-- dimos/memory2/cli/render.py | 4 +-- dimos/memory2/cli/summary.py | 36 ++++++++----------- dimos/memory2/codecs/base.py | 4 +-- dimos/memory2/store/sqlite.py | 4 +-- dimos/memory2/stream.py | 7 ++-- dimos/memory2/utils/progress.py | 1 + dimos/protocol/pubsub/benchmark/type.py | 25 ++----------- dimos/utils/colors.py | 5 +++ 11 files changed, 38 insertions(+), 66 deletions(-) diff --git a/dimos/mapping/utils/cli/dataset_validation.md b/dimos/mapping/utils/cli/dataset_validation.md index 9e4873ed91..79b8631b54 100644 --- a/dimos/mapping/utils/cli/dataset_validation.md +++ b/dimos/mapping/utils/cli/dataset_validation.md @@ -2,15 +2,11 @@ Dataset Validation ```sh dimos mem summary recording_go2_mid360_2026-05-29_4-45pm-PST.db - -Stream("color_image"): 11141 items, 2026-05-29 23:32:57 — 2026-05-29 23:45:57 (780.1s) -Stream("fastlio_lidar"): 7240 items, 2026-05-29 23:32:56 — 2026-05-29 23:45:57 (781.7s) -Stream("fastlio_odometry"): 18737 items, 2026-05-29 23:32:56 — 2026-05-29 23:45:57 (781.8s) -Stream("lidar"): 6025 items, 2026-05-29 23:32:55 — 2026-05-29 23:45:57 (782.3s) -Stream("odom"): 14630 items, 2026-05-29 23:32:55 — 2026-05-29 23:45:57 (782.3s) ``` -Shows which streams are in the database. You can replay messages in rerun: +Prints a table of the streams in the database: per-stream item count, average +Hz, start time (UTC), duration, and stored size, plus a total row. You can +replay messages in rerun: ```sh dimos map replay recording_go2_mid360_2026-05-29_4-45pm-PST.db --duration 60 diff --git a/dimos/mapping/utils/cli/rename.py b/dimos/mapping/utils/cli/rename.py index e856131e86..a55701ccfb 100644 --- a/dimos/mapping/utils/cli/rename.py +++ b/dimos/mapping/utils/cli/rename.py @@ -31,7 +31,7 @@ import typer -from dimos.memory2.codecs.base import _resolve_payload_type +from dimos.memory2.codecs.base import resolve_payload_type from dimos.memory2.store.sqlite import SqliteStore from dimos.memory2.stream import Stream from dimos.memory2.type.observation import Observation @@ -77,7 +77,7 @@ def _stream_payload_types(db_path: Path) -> dict[str, type]: rows = conn.execute("SELECT name, config FROM _streams").fetchall() finally: conn.close() - return {name: _resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} + return {name: resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} def _parse_renames(pairs: list[str]) -> dict[str, str]: diff --git a/dimos/mapping/utils/cli/replay.py b/dimos/mapping/utils/cli/replay.py index b056f44321..e29d2b02b1 100644 --- a/dimos/mapping/utils/cli/replay.py +++ b/dimos/mapping/utils/cli/replay.py @@ -175,7 +175,7 @@ def main( ) -> None: """Dump a recording to .rrd (lidar clouds + camera frames) and open it in rerun.""" from dimos.mapping.voxels import VoxelMapTransformer - from dimos.memory2.cli.summary import _stream_payload_types + from dimos.memory2.cli.summary import stream_payload_types from dimos.memory2.store.sqlite import SqliteStore from dimos.memory2.transform import throttle from dimos.msgs.geometry_msgs.PoseStamped import PoseStamped @@ -192,7 +192,7 @@ def main( # Resolve which streams to voxelize: all PointCloud2 streams, or the # explicit --map-source subset. Validate up front so typos fail fast. - pc_streams = [n for n, t in _stream_payload_types(db_path).items() if t is PointCloud2] + pc_streams = [n for n, t in stream_payload_types(db_path).items() if t is PointCloud2] map_sources = list(map_source) or pc_streams if (map or map_final) and (bad := [s for s in map_sources if s not in pc_streams]): raise typer.BadParameter(f"--map-source: not PointCloud2 stream(s): {', '.join(bad)}") diff --git a/dimos/memory2/cli/render.py b/dimos/memory2/cli/render.py index 615fa77546..1f816f7c78 100644 --- a/dimos/memory2/cli/render.py +++ b/dimos/memory2/cli/render.py @@ -28,6 +28,8 @@ import subprocess from typing import TYPE_CHECKING +from dimos.utils.data import resolve_named_path + if TYPE_CHECKING: from dimos.memory2.store.base import Store @@ -38,8 +40,6 @@ def open_store(path: str) -> Store: Bare names resolve like the other db verbs: cwd, then ``data/``, then an LFS pull — defaulting to ``.db`` unless the name says ``.mcap``. """ - from dimos.utils.data import resolve_named_path - resolved = resolve_named_path(path, ".mcap" if str(path).endswith(".mcap") else ".db") if resolved.suffix == ".mcap": from dimos.robot.unitree.go2.dds.store import Go2McapStore # lazy: robot-layer codec set diff --git a/dimos/memory2/cli/summary.py b/dimos/memory2/cli/summary.py index 95b111ce3f..8aa2519915 100644 --- a/dimos/memory2/cli/summary.py +++ b/dimos/memory2/cli/summary.py @@ -20,48 +20,46 @@ from __future__ import annotations +from datetime import datetime, timezone +import json +from math import log +import sqlite3 from typing import TYPE_CHECKING, Any import typer +from dimos.utils.colors import HEAT_GRADIENT_ANSI256 +from dimos.utils.data import resolve_named_path +from dimos.utils.human import human_bytes + if TYPE_CHECKING: from pathlib import Path from dimos.memory2.stream import Stream -# Heavy dimos imports (memory2 store → codecs, msgs) are deferred into the -# function bodies so that `dimos --help` — which imports this module just to +# Heavy dimos imports (memory2 store → codecs, msgs) and rich are deferred into +# the function bodies so that `dimos --help` — which imports this module just to # register the `mem summary` command — stays fast. See test_cli_startup.py. -def _stream_payload_types(db_path: Path) -> dict[str, type]: +def stream_payload_types(db_path: Path) -> dict[str, type]: """Read each stream's registered payload type from the _streams table.""" - import json - import sqlite3 - - from dimos.memory2.codecs.base import _resolve_payload_type + from dimos.memory2.codecs.base import resolve_payload_type conn = sqlite3.connect(str(db_path)) try: rows = conn.execute("SELECT name, config FROM _streams").fetchall() finally: conn.close() - return {name: _resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} - - -# ANSI 256 gradient: red -> orange -> yellow -> green (same as the pubsub -# benchmark heatmaps in dimos/protocol/pubsub/benchmark/type.py). -_GRADIENT = [52, 88, 124, 160, 196, 202, 208, 214, 220, 226, 190, 154, 148, 118, 82, 46, 40, 34] + return {name: resolve_payload_type(json.loads(cfg)["payload_module"]) for name, cfg in rows} def _shade(value: float, lo: float, hi: float) -> str: """Rich style for ``value`` relative to [lo, hi], log-scaled (columns span decades).""" - from math import log - if value <= 0: return "dim" t = 0.5 if hi <= lo else (log(value) - log(lo)) / (log(hi) - log(lo)) - return f"color({_GRADIENT[round(t * (len(_GRADIENT) - 1))]})" + return f"color({HEAT_GRADIENT_ANSI256[round(t * (len(HEAT_GRADIENT_ANSI256) - 1))]})" def _heat(text: str, value: float, column: list[float]) -> str: @@ -75,18 +73,14 @@ def main( dataset: str = typer.Argument(..., help="Dataset .db: bare name (cwd or data/) or path"), ) -> None: """Print per-stream counts, time ranges, and payload sizes for a recorded SQLite dataset.""" - from datetime import datetime, timezone - from rich.console import Console from rich.progress import Progress from rich.table import Table from dimos.memory2.store.sqlite import SqliteStore - from dimos.utils.data import resolve_named_path - from dimos.utils.human import human_bytes db_path = resolve_named_path(dataset, ".db") - payload_types = _stream_payload_types(db_path) + payload_types = stream_payload_types(db_path) rows: list[tuple[str, int, float | None, float | None, int]] = [] store = SqliteStore(path=str(db_path)) diff --git a/dimos/memory2/codecs/base.py b/dimos/memory2/codecs/base.py index 821b36b60f..0ddd384e84 100644 --- a/dimos/memory2/codecs/base.py +++ b/dimos/memory2/codecs/base.py @@ -81,7 +81,7 @@ def _class_to_id(codec: Any) -> str: return name.lower() -def _resolve_payload_type(payload_module: str) -> type[Any]: +def resolve_payload_type(payload_module: str) -> type[Any]: parts = payload_module.rsplit(".", 1) if len(parts) != 2: raise ValueError(f"Cannot resolve payload type from {payload_module!r}") @@ -104,7 +104,7 @@ def _make_one(name: str, payload_module: str, inner: Codec[Any] | None = None) - if name == "lcm": from dimos.memory2.codecs.lcm import LcmCodec - return LcmCodec(_resolve_payload_type(payload_module)) + return LcmCodec(resolve_payload_type(payload_module)) if name == "pickle": from dimos.memory2.codecs.pickle import PickleCodec diff --git a/dimos/memory2/store/sqlite.py b/dimos/memory2/store/sqlite.py index 229961a126..8afab3a714 100644 --- a/dimos/memory2/store/sqlite.py +++ b/dimos/memory2/store/sqlite.py @@ -69,11 +69,11 @@ def _open_connection(self) -> sqlite3.Connection: def _assemble_backend(self, name: str, stored: dict[str, Any]) -> Backend[Any]: """Reconstruct a Backend from a stored config dict.""" - from dimos.memory2.codecs.base import _resolve_payload_type, codec_from_id + from dimos.memory2.codecs.base import codec_from_id, resolve_payload_type payload_module = stored["payload_module"] codec = codec_from_id(stored["codec_id"], payload_module) - data_type = _resolve_payload_type(payload_module) + data_type = resolve_payload_type(payload_module) eager_blobs = stored.get("eager_blobs", False) page_size = stored.get("page_size", self.config.page_size) diff --git a/dimos/memory2/stream.py b/dimos/memory2/stream.py index 976e47afef..539a002b13 100644 --- a/dimos/memory2/stream.py +++ b/dimos/memory2/stream.py @@ -26,6 +26,7 @@ from typing_extensions import TypeVar from dimos.core.resource import CompositeResource +from dimos.memory2.backend import Backend from dimos.memory2.buffer import BackpressureBuffer, KeepLast from dimos.memory2.transform import FnIterTransformer, FnTransformer, Transformer from dimos.memory2.type.filter import ( @@ -41,6 +42,7 @@ ) from dimos.memory2.type.observation import EmbeddedObservation, Observation from dimos.msgs.geometry_msgs.Vector3 import Vector3 +from dimos.utils.human import human_bytes from dimos.utils.logging_config import setup_logger if TYPE_CHECKING: @@ -49,7 +51,6 @@ import reactivex from reactivex.abc import DisposableBase, ObserverBase - from dimos.memory2.backend import Backend from dimos.models.embedding.base import Embedding T = TypeVar("T") @@ -498,8 +499,6 @@ def get_time_range(self) -> tuple[float, float]: def size_bytes(self) -> int | None: """Total stored payload bytes, or None if unknown or the query narrows the stream.""" - from dimos.memory2.backend import Backend - q = self._query narrowed = ( q.filters @@ -518,8 +517,6 @@ def summary(self) -> str: """Return a short human-readable summary: count, time range, avg frequency, size.""" from datetime import datetime, timezone - from dimos.utils.human import human_bytes - sz = self.size_bytes() size = f", {human_bytes(sz)}" if sz is not None else "" diff --git a/dimos/memory2/utils/progress.py b/dimos/memory2/utils/progress.py index cf8cea5b0c..23504d7896 100644 --- a/dimos/memory2/utils/progress.py +++ b/dimos/memory2/utils/progress.py @@ -50,6 +50,7 @@ class ProgressBar: """Callable per-observation progress; ``close()`` finalizes an unfinished bar.""" def __init__(self, total: int, label: str = "") -> None: + # lazy: keep rich off the CLI startup path (see test_cli_startup.py) from rich.progress import ( BarColumn, MofNCompleteColumn, diff --git a/dimos/protocol/pubsub/benchmark/type.py b/dimos/protocol/pubsub/benchmark/type.py index eea43ac44c..6fb717e983 100644 --- a/dimos/protocol/pubsub/benchmark/type.py +++ b/dimos/protocol/pubsub/benchmark/type.py @@ -20,6 +20,7 @@ from typing import Any, Generic from dimos.protocol.pubsub.spec import MsgT, PubSub, TopicT +from dimos.utils.colors import HEAT_GRADIENT_ANSI256 from dimos.utils.human import human_bytes, human_duration, human_number MsgGen = Callable[[int], tuple[TopicT, MsgT]] @@ -164,29 +165,7 @@ def _print_heatmap( return min_val, max_val = min(all_vals), max(all_vals) - # ANSI 256 gradient: red -> orange -> yellow -> green - gradient = [ - 52, - 88, - 124, - 160, - 196, - 202, - 208, - 214, - 220, - 226, - 190, - 154, - 148, - 118, - 82, - 46, - 40, - 34, - ] - if not high_is_good: - gradient = gradient[::-1] + gradient = HEAT_GRADIENT_ANSI256 if high_is_good else HEAT_GRADIENT_ANSI256[::-1] def val_to_color(v: float) -> int: if v <= 0 or max_val == min_val: diff --git a/dimos/utils/colors.py b/dimos/utils/colors.py index 294cf5d43b..8f5507ddfc 100644 --- a/dimos/utils/colors.py +++ b/dimos/utils/colors.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +# ANSI 256 heatmap gradient: red -> orange -> yellow -> green. +HEAT_GRADIENT_ANSI256: list[int] = [ + 52, 88, 124, 160, 196, 202, 208, 214, 220, 226, 190, 154, 148, 118, 82, 46, 40, 34, +] # fmt: skip + def green(text: str) -> str: """Return the given text in green color.""" From 5a432db65edb4eb5c4a27fb079c1dbcfb98eeb4e Mon Sep 17 00:00:00 2001 From: "brain (osvetnik)" Date: Tue, 7 Jul 2026 01:49:42 +0300 Subject: [PATCH 11/11] Make progress() a context manager; drop module-global cleanup Per review: progress() now yields the bar via contextlib.contextmanager and finalizes it on exit (including on exception or early break), so the _current global and the atexit hook are gone. All call sites use 'with progress(...) as bar'. Adds tests for cleanup-on-exception, early-exit, and post-completion calls. --- dimos/mapping/utils/cli/map.py | 78 +++++++------- dimos/mapping/utils/cli/replay.py | 124 +++++++++++----------- dimos/memory2/cli/render.py | 29 +++-- dimos/memory2/utils/progress.py | 48 ++++----- dimos/memory2/utils/test_progress.py | 75 +++++++++++++ dimos/robot/unitree/go2/dds/cli/render.py | 59 +++++----- 6 files changed, 246 insertions(+), 167 deletions(-) create mode 100644 dimos/memory2/utils/test_progress.py diff --git a/dimos/mapping/utils/cli/map.py b/dimos/mapping/utils/cli/map.py index 106e141797..de003bf15e 100644 --- a/dimos/mapping/utils/cli/map.py +++ b/dimos/mapping/utils/cli/map.py @@ -526,50 +526,53 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: graph: PoseGraph | None = None if pgo: print("running PGO twopass map...") - prog = progress(total, "pgo pass 1 (optimizing)") - graph = lidar.tap(prog).transform(PGO()).last().data + with progress(total, "pgo pass 1 (optimizing)") as bar: + graph = lidar.tap(bar).transform(PGO()).last().data pgo_path = [ (kf.optimized.translation.x, kf.optimized.translation.y, kf.optimized.translation.z) for kf in graph.keyframes ] - pgo_map = _accumulate( - kept, - voxel=voxel, - block_count=block_count, - device=device, - graph=graph, - register=register, - carve_columns=carve, - progress_cb=progress(n_kept, "pgo pass 2 (rebuilding)"), - ) + with progress(n_kept, "pgo pass 2 (rebuilding)") as bar: + pgo_map = _accumulate( + kept, + voxel=voxel, + block_count=block_count, + device=device, + graph=graph, + register=register, + carve_columns=carve, + progress_cb=bar, + ) full_pgo_map = None if full_pgo: assert graph is not None - full_pgo_map = _accumulate( - lidar, + with progress(total, "full pgo (rebuilding)") as bar: + full_pgo_map = _accumulate( + lidar, + voxel=voxel, + block_count=block_count, + device=device, + graph=graph, + register=register, + carve_columns=carve, + progress_cb=bar, + ) + + # Raw map: same dedup'd frames, no PGO correction. + with progress(n_kept, "reconstructing global map") as bar: + global_map = _accumulate( + kept, voxel=voxel, block_count=block_count, device=device, - graph=graph, register=register, carve_columns=carve, - progress_cb=progress(total, "full pgo (rebuilding)"), + progress_cb=bar, ) - # Raw map: same dedup'd frames, no PGO correction. - global_map = _accumulate( - kept, - voxel=voxel, - block_count=block_count, - device=device, - register=register, - carve_columns=carve, - progress_cb=progress(n_kept, "reconstructing global map"), - ) - marker_dets: list[Observation[Any]] = [] if markers: # Image observations in dimos recordings are stamped with @@ -596,17 +599,18 @@ def _position(obs: Observation[Any]) -> tuple[float, float, float] | None: # Keep the sharpest frame per --marker-quality-window window, then # drop frames where the robot was moving (linear + rotational) faster # than the limits. Defaults match replay_marker.py so positions agree. - pipeline: Stream[Image] = color_image.tap( - progress(n_images, "detecting markers") - ).transform(QualityWindow(lambda img: img.sharpness, window=marker_quality_window)) - if marker_max_speed > 0: - pipeline = pipeline.transform( - SpeedLimit( - max_mps=marker_max_speed, - max_dps=marker_max_rot_rate if marker_max_rot_rate > 0 else None, - ) + with progress(n_images, "detecting markers") as bar: + pipeline: Stream[Image] = color_image.tap(bar).transform( + QualityWindow(lambda img: img.sharpness, window=marker_quality_window) ) - all_dets = pipeline.transform(xf).to_list() + if marker_max_speed > 0: + pipeline = pipeline.transform( + SpeedLimit( + max_mps=marker_max_speed, + max_dps=marker_max_rot_rate if marker_max_rot_rate > 0 else None, + ) + ) + all_dets = pipeline.transform(xf).to_list() if marker_smoothing > 0: # Keep only the latest emission per track_id — that's the most # averaged pose, drawn once per tracked marker session. diff --git a/dimos/mapping/utils/cli/replay.py b/dimos/mapping/utils/cli/replay.py index e29d2b02b1..209ab6ee95 100644 --- a/dimos/mapping/utils/cli/replay.py +++ b/dimos/mapping/utils/cli/replay.py @@ -67,14 +67,14 @@ def _log_clouds( whole pipeline. """ n = total if total is not None else stream.count() - cb = progress(n, label) - for obs in stream: - cb(obs) - rr.set_time(TIMELINE, timestamp=obs.ts) - rr.log( - entity, - obs.data.to_rerun(voxel_size=voxel, mode=point_mode, bottom_cutoff=bottom_cutoff), - ) + with progress(n, label) as bar: + for obs in stream: + bar(obs) + rr.set_time(TIMELINE, timestamp=obs.ts) + rr.log( + entity, + obs.data.to_rerun(voxel_size=voxel, mode=point_mode, bottom_cutoff=bottom_cutoff), + ) def _log_path( @@ -90,22 +90,22 @@ def _log_path( without a pose are skipped. """ n = stream.count() - cb = progress(n, label) points: list[tuple[float, float, float]] = [] last_ts: float | None = None emit_count = 0 - for obs in stream: - cb(obs) - if obs.pose_tuple is None: - continue - points.append( - (float(obs.pose_tuple[0]), float(obs.pose_tuple[1]), float(obs.pose_tuple[2])) - ) - last_ts = obs.ts - emit_count += 1 - if emit_every > 0 and emit_count % emit_every == 0 and len(points) >= 2: - rr.set_time(TIMELINE, timestamp=obs.ts) - rr.log(entity, rr.LineStrips3D([points], colors=[color])) + with progress(n, label) as bar: + for obs in stream: + bar(obs) + if obs.pose_tuple is None: + continue + points.append( + (float(obs.pose_tuple[0]), float(obs.pose_tuple[1]), float(obs.pose_tuple[2])) + ) + last_ts = obs.ts + emit_count += 1 + if emit_every > 0 and emit_count % emit_every == 0 and len(points) >= 2: + rr.set_time(TIMELINE, timestamp=obs.ts) + rr.log(entity, rr.LineStrips3D([points], colors=[color])) if ( last_ts is not None and len(points) >= 2 @@ -277,20 +277,20 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: # fastlio pose axis + path from fastlio_odometry stream. if "fastlio_odometry" in store.streams: odometry = clipped("fastlio_odometry", Odometry) - cb = progress(odometry.count(), "fastlio_odometry") - for obs in odometry: - cb(obs) - if obs.pose_tuple is None: - continue - rr.set_time(TIMELINE, timestamp=obs.ts) - x, y, z, qx, qy, qz, qw = obs.pose_tuple - rr.log( - "world/fastlio", - rr.Transform3D( - translation=[x, y, z], - quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), - ), - ) + with progress(odometry.count(), "fastlio_odometry") as bar: + for obs in odometry: + bar(obs) + if obs.pose_tuple is None: + continue + rr.set_time(TIMELINE, timestamp=obs.ts) + x, y, z, qx, qy, qz, qw = obs.pose_tuple + rr.log( + "world/fastlio", + rr.Transform3D( + translation=[x, y, z], + quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), + ), + ) _log_path( " fastlio_path", clipped("fastlio_odometry", Odometry), @@ -301,20 +301,20 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: # Go2 native odom pose axis + path. if "odom" in store.streams: odom = clipped("odom", PoseStamped) - cb = progress(odom.count(), " odom") - for odom_obs in odom: - cb(odom_obs) - if odom_obs.pose_tuple is None: - continue - rr.set_time(TIMELINE, timestamp=odom_obs.ts) - x, y, z, qx, qy, qz, qw = odom_obs.pose_tuple - rr.log( - "world/odom", - rr.Transform3D( - translation=[x, y, z], - quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), - ), - ) + with progress(odom.count(), " odom") as bar: + for odom_obs in odom: + bar(odom_obs) + if odom_obs.pose_tuple is None: + continue + rr.set_time(TIMELINE, timestamp=odom_obs.ts) + x, y, z, qx, qy, qz, qw = odom_obs.pose_tuple + rr.log( + "world/odom", + rr.Transform3D( + translation=[x, y, z], + quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]), + ), + ) _log_path( " odom_path", clipped("odom", PoseStamped), @@ -327,19 +327,19 @@ def clipped(name: str, ptype: type[Any]) -> Stream[Any]: color_image.transform(throttle(1.0 / camera_hz)) if camera_hz > 0 else color_image ) n_img = cam_pipeline.count() - cb = progress(n_img, " color_image") - for img_obs in cam_pipeline: - cb(img_obs) - rr.set_time(TIMELINE, timestamp=img_obs.ts) - if img_obs.pose_tuple is not None: - x, y, z, qx, qy, qz, qw = img_obs.pose_tuple - rr.log( - "world/camera", - rr.Transform3D( - translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) - ), - ) - rr.log("world/camera/image", img_obs.data.to_rerun()) + with progress(n_img, " color_image") as bar: + for img_obs in cam_pipeline: + bar(img_obs) + rr.set_time(TIMELINE, timestamp=img_obs.ts) + if img_obs.pose_tuple is not None: + x, y, z, qx, qy, qz, qw = img_obs.pose_tuple + rr.log( + "world/camera", + rr.Transform3D( + translation=[x, y, z], quaternion=rr.Quaternion(xyzw=[qx, qy, qz, qw]) + ), + ) + rr.log("world/camera/image", img_obs.data.to_rerun()) print(f"wrote {out}") if no_gui: diff --git a/dimos/memory2/cli/render.py b/dimos/memory2/cli/render.py index 1f816f7c78..fe50c0d489 100644 --- a/dimos/memory2/cli/render.py +++ b/dimos/memory2/cli/render.py @@ -109,22 +109,21 @@ def render_store( rr.save(out) for name, stream in renderable: - report = progress(stream.count(), label=name) - for obs in stream: - if seconds is not None and obs.ts - t0 > seconds: - report.close() # finalize the windowed (sub-100%) bar - break - if obs.data is None: # e.g. a truncated/corrupt frame that failed to decode + with progress(stream.count(), label=name) as report: + for obs in stream: + if seconds is not None and obs.ts - t0 > seconds: + break # the context manager finalizes the windowed (sub-100%) bar + if obs.data is None: # e.g. a truncated/corrupt frame that failed to decode + report(obs) + continue + rr.set_time("time", duration=obs.ts - t0) + data = obs.data.to_rerun() + if isinstance(data, list): # RerunMulti: [(subpath, archetype), ...] + for sub, arch in data: + rr.log(f"{name}/{sub}", arch) + else: + rr.log(name, data) report(obs) - continue - rr.set_time("time", duration=obs.ts - t0) - data = obs.data.to_rerun() - if isinstance(data, list): # RerunMulti: [(subpath, archetype), ...] - for sub, arch in data: - rr.log(f"{name}/{sub}", arch) - else: - rr.log(name, data) - report(obs) rr.rerun_shutdown() # flush + close the .rrd before opening it print(f"wrote {out}") diff --git a/dimos/memory2/utils/progress.py b/dimos/memory2/utils/progress.py index 23504d7896..05f1256ac1 100644 --- a/dimos/memory2/utils/progress.py +++ b/dimos/memory2/utils/progress.py @@ -14,37 +14,31 @@ """A per-observation progress bar, shaped for ``Stream.tap``. - stream.tap(progress(stream.count(), "render")).drain() + with progress(stream.count(), "render") as bar: + stream.tap(bar).drain() Renders a rich progress bar: label, percent, count, data-seconds covered, speed relative to wall-clock (``x rt``), and per-frame latency. On completion -the live bar is replaced by one plain persisted line. One bar is live at a -time — starting a new bar finalizes a dangling one (rich allows a single live -display); call :meth:`ProgressBar.close` when abandoning a stream early -(e.g. a bounded time window). +the live bar is replaced by one plain persisted line. ``progress`` is a +context manager: exiting the ``with`` block finalizes the bar (and restores +the cursor), even when the stream is abandoned early (e.g. a bounded time +window) or the pipeline raises. rich allows a single live display, so keep +one bar per ``with`` block at a time. """ from __future__ import annotations -import atexit +from contextlib import contextmanager import time from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + from collections.abc import Iterator + from dimos.memory2.type.observation import Observation _REFRESH_S = 0.1 # refresh cap; keeps 200 Hz streams from redrawing per frame -_current: ProgressBar | None = None - - -def _close_current() -> None: - if _current is not None: - _current.close() - - -atexit.register(_close_current) # restore the cursor if a bar is abandoned - class ProgressBar: """Callable per-observation progress; ``close()`` finalizes an unfinished bar.""" @@ -86,10 +80,6 @@ def __call__(self, obs: Observation[Any]) -> None: return # observations may keep flowing after close (e.g. a tap downstream) now = time.monotonic() if self._wall_start is None: - global _current - if _current is not None: - _current.close() # only one live display may be active - _current = self self._wall_start = now self._first_ts = obs.ts self._prog.start() @@ -110,17 +100,23 @@ def __call__(self, obs: Observation[Any]) -> None: def close(self) -> None: """Clear the live bar and persist it as a plain line.""" - global _current if self._closed: return self._closed = True - if _current is self: - _current = None self._prog.stop() # transient: wipes the live bar pct = 100 * self._seen // self._total if self._total else 100 print(f"{self._label} {pct}% [{self._seen}/{self._total}] {self._stats}") -def progress(total: int, label: str = "") -> ProgressBar: - """Return a callback that renders streaming progress, one call per observation.""" - return ProgressBar(total, label) +@contextmanager +def progress(total: int, label: str = "") -> Iterator[ProgressBar]: + """Yield a per-observation progress callback; the bar finalizes on exit. + + with progress(stream.count(), "render") as bar: + stream.tap(bar).drain() + """ + bar = ProgressBar(total, label) + try: + yield bar + finally: + bar.close() diff --git a/dimos/memory2/utils/test_progress.py b/dimos/memory2/utils/test_progress.py new file mode 100644 index 0000000000..309ca71408 --- /dev/null +++ b/dimos/memory2/utils/test_progress.py @@ -0,0 +1,75 @@ +# Copyright 2026 Dimensional Inc. +# +# 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. + +"""Tests for the ``progress`` context manager.""" + +from __future__ import annotations + +import pytest + +from dimos.memory2.type.observation import Observation +from dimos.memory2.utils.progress import progress + + +def _obs(ts: float) -> Observation[None]: + return Observation(ts=ts, _data=None) + + +def test_completed_bar_persists_one_line(capsys: pytest.CaptureFixture[str]) -> None: + with progress(3, "render") as bar: + for i in range(3): + bar(_obs(float(i))) + out = capsys.readouterr().out + assert out.count("render 100% [3/3]") == 1 + + +def test_early_exit_finalizes_partial_bar(capsys: pytest.CaptureFixture[str]) -> None: + with progress(4, "windowed") as bar: + bar(_obs(0.0)) + bar(_obs(0.1)) + out = capsys.readouterr().out + assert out.count("windowed 50% [2/4]") == 1 + + +def test_cleanup_on_exception(capsys: pytest.CaptureFixture[str]) -> None: + with pytest.raises(RuntimeError, match="boom"): + with progress(10, "crashy") as bar: + bar(_obs(0.0)) + raise RuntimeError("boom") + out = capsys.readouterr().out + assert out.count("crashy 10% [1/10]") == 1 + + +def test_calls_after_completion_are_ignored(capsys: pytest.CaptureFixture[str]) -> None: + with progress(2, "tap") as bar: + for i in range(5): # observations may keep flowing after the bar completes + bar(_obs(float(i))) + out = capsys.readouterr().out + assert out.count("tap 100% [2/2]") == 1 + + +def test_zero_total_reports_complete(capsys: pytest.CaptureFixture[str]) -> None: + with progress(0, "empty"): + pass + out = capsys.readouterr().out + assert out.count("empty 100% [0/0]") == 1 + + +def test_sequential_bars_each_persist(capsys: pytest.CaptureFixture[str]) -> None: + for label in ("first", "second"): + with progress(1, label) as bar: + bar(_obs(0.0)) + out = capsys.readouterr().out + assert out.count("first 100% [1/1]") == 1 + assert out.count("second 100% [1/1]") == 1 diff --git a/dimos/robot/unitree/go2/dds/cli/render.py b/dimos/robot/unitree/go2/dds/cli/render.py index ea494b41ef..748db33933 100644 --- a/dimos/robot/unitree/go2/dds/cli/render.py +++ b/dimos/robot/unitree/go2/dds/cli/render.py @@ -167,15 +167,16 @@ def log_path(obs: Observation[Path]) -> None: rr.log("world/leg_odom_path", obs.data.to_rerun()) src = store.streams.sportmodestate.to_time(seconds) - ( - src.tap(progress(src.count(), "leg_odom")) - .map_data(sportmode_pose) - .tap(log_pose) - .transform(throttle(0.1)) # reduce_rate: thin the path to ~10 Hz - .transform(accumulate_path) # yield the growing path each step - .tap(log_path) - .drain() - ) + with progress(src.count(), "leg_odom") as bar: + ( + src.tap(bar) + .map_data(sportmode_pose) + .tap(log_pose) + .transform(throttle(0.1)) # reduce_rate: thin the path to ~10 Hz + .transform(accumulate_path) # yield the growing path each step + .tap(log_path) + .drain() + ) def imu_odom(store: Go2McapStore, seconds: float | None) -> None: @@ -187,15 +188,16 @@ def log_path(obs: Observation[Path]) -> None: gravity = gravity_bias(store) # calibrate gravity+bias on the stationary start src = store.streams.imu.to_time(seconds) - ( - src.tap(progress(src.count(), "imu_odom")) - .scan_data((np.zeros(3), None), integrate_velocity(gravity)) # -> velocity (TwistStamped) - .scan_data((np.zeros(3), None), integrate_position) # -> position (PoseStamped) - .transform(throttle(0.1)) # thin the path after integrating at full IMU rate - .transform(accumulate_path) - .tap(log_path) - .drain() - ) + with progress(src.count(), "imu_odom") as bar: + ( + src.tap(bar) + .scan_data((np.zeros(3), None), integrate_velocity(gravity)) # -> velocity + .scan_data((np.zeros(3), None), integrate_position) # -> position (PoseStamped) + .transform(throttle(0.1)) # thin the path after integrating at full IMU rate + .transform(accumulate_path) + .tap(log_path) + .drain() + ) def lidar(store: Go2McapStore, seconds: float | None) -> None: @@ -207,7 +209,8 @@ def log_lidar(obs: Observation[PointCloud2]) -> None: src = store.streams.lidar.to_time(seconds) rr.log("world/leg_odom/lidar", LIDAR_TO_BASE.to_rerun(frameless=True), static=True) - (src.tap(progress(src.count(), "lidar")).tap(log_lidar).drain()) + with progress(src.count(), "lidar") as bar: + src.tap(bar).tap(log_lidar).drain() def _interp_pose( @@ -252,13 +255,14 @@ def log_voxels(obs: Observation[PointCloud2]) -> None: rr.log("world/world_lidar", obs.data.to_rerun()) src = store.streams.lidar.to_time(seconds) - ( - src.tap(progress(src.count(), "world_lidar")) - .map_data(to_world) # lidar cloud -> world-frame cloud - .transform(VoxelMapTransformer(emit_every=10, voxel_size=0.1)) # global voxel map - .tap(log_voxels) - .drain() - ) + with progress(src.count(), "world_lidar") as bar: + ( + src.tap(bar) + .map_data(to_world) # lidar cloud -> world-frame cloud + .transform(VoxelMapTransformer(emit_every=10, voxel_size=0.1)) # global voxel map + .tap(log_voxels) + .drain() + ) def camera(store: Go2McapStore, seconds: float | None, hz: float) -> None: @@ -275,7 +279,8 @@ def log_image(obs: Observation[Image]) -> None: rr.log("world/camera", obs.data.to_rerun()) src = store.streams.color_image.to_time(seconds) - (src.tap(progress(src.count(), "camera")).transform(throttle(1.0 / hz)).tap(log_image).drain()) + with progress(src.count(), "camera") as bar: + src.tap(bar).transform(throttle(1.0 / hz)).tap(log_image).drain() # Add a source: write a (store, seconds) -> None function and append it.