Fix non terminating speed/elapsed in marimo progress bar

Author: ricardoV94

pymc/progress_bar/marimo_progress.py

@@ -71,6 +71,7 @@ def __init__(
         self._mo_replace: Callable[[object], None] | None = None
         self._task_state: list[dict[str, Any]] = []
         self._start_times: list[float | None] = []
+        self._end_times: list[float | None] = []
         self._last_render_time: float = 0.0
         self._min_render_interval: float = 0.1
 
@@ -107,6 +108,7 @@ def _initialize_tasks(self) -> None:
         # ``cores < chains``, later chains aren't timed from the first chain's
         # start.
         self._start_times = [None] * self.n_bars
+        self._end_times = [None] * self.n_bars
 
     def update(
         self,
@@ -148,6 +150,8 @@ def update(
 
         if is_last:
             self._task_state[task_id]["completed"] = self._task_state[task_id]["total"]
+            if self._end_times[task_id] is None:
+                self._end_times[task_id] = perf_counter()
 
         now = perf_counter()
         if is_last or (now - self._last_render_time) >= self._min_render_interval:
@@ -215,7 +219,12 @@ def _render_task_row(self, task_id: int, state: dict[str, Any], stat_keys: list[
 
         pct = (completed / total * 100) if total else 0
         start_time = self._start_times[task_id]
-        elapsed = 0.0 if start_time is None else perf_counter() - start_time
+        if start_time is None:
+            elapsed = 0.0
+        else:
+            end_time = self._end_times[task_id]
+            now = end_time if end_time is not None else perf_counter()
+            elapsed = now - start_time
 
         action = self.step_name.lower()
         # Wait for a small window of elapsed time before computing speed so the

pymc/progress_bar/progress.py

@@ -372,6 +372,7 @@ def __init__(
         )
         # Used to compute delta draws between calls
         self._previous_finished = [0] * chains
+        self._chain_completed = [False] * chains
 
         progress_columns = [
             TextColumn("{task.fields[divergences]}", table_column=Column("Divergences", ratio=1)),
@@ -424,11 +425,14 @@ def update(self, chain_progresses) -> None:
             for chain_idx, cp in enumerate(chain_progresses):
                 # With ``cores < chains`` queued chains haven't started yet;
                 # skip them so their bar doesn't show progress or elapsed time.
-                if not cp.started:
+                if not cp.started or self._chain_completed[chain_idx]:
                     continue
                 cp_finished_draws = cp.finished_draws
                 delta = cp_finished_draws - self._previous_finished[chain_idx]
                 self._previous_finished[chain_idx] = cp_finished_draws
+                is_last = cp_finished_draws >= cp.total_draws
+                if is_last:
+                    self._chain_completed[chain_idx] = True
                 # Use nutpie's per-chain runtime as the source of truth for
                 # elapsed/speed, so reads aren't skewed by the wait time
                 # before this chain started.
@@ -445,7 +449,7 @@ def update(self, chain_progresses) -> None:
                     advance=delta,
                     failing=bool(cp.divergent_draws),
                     stats=stats,
-                    is_last=cp_finished_draws >= cp.total_draws,
+                    is_last=is_last,
                     total=cp.total_draws,
                 )
 

tests/progress_bar/test_marimo.py

@@ -12,13 +12,15 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+from time import sleep
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import pytest
 
 import pymc as pm
 
-from pymc.progress_bar import MCMCProgressBarManager
+from pymc.progress_bar import MCMCProgressBarManager, NutpieProgressBarManager
 from pymc.progress_bar.marimo_progress import MarimoProgressBackend
 
 
@@ -75,6 +77,7 @@ def test_render_html_structure(self, step_method):
                 {"completed": 75, "total": 150, "failing": True, "stats": {"divergences": 1}},
             ]
             backend._start_times = [0, 0]
+            backend._end_times = [None, None]
 
             html = backend._render_html()
 
@@ -103,6 +106,7 @@ def test_render_html_with_stats(self, step_method):
                 },
             ]
             backend._start_times = [0]
+            backend._end_times = [None]
 
             html = backend._render_html()
 
@@ -121,6 +125,25 @@ def test_is_last_sets_completed_to_total(self):
         assert backend._task_state[0]["completed"] == 150
         assert backend._task_state[1]["completed"] == 0
 
+    def test_elapsed_freezes_after_completion(self):
+        """Completed chains must not show drifting speed/elapsed on re-renders."""
+        backend = MarimoProgressBackend(
+            step_name="Draw", n_bars=2, total=10, combined=False, full_stats=False
+        )
+        backend._initialize_tasks()
+
+        # Complete chain 0
+        for i in range(10):
+            backend.update(task_id=0, advance=1, failing=False, stats={}, is_last=i == 9)
+
+        html_at_finish = backend._render_task_row(0, backend._task_state[0], [])
+
+        # Let wall-clock advance so unfrozen elapsed would visibly drift
+        sleep(0.3)
+
+        html_after = backend._render_task_row(0, backend._task_state[0], [])
+        assert html_at_finish == html_after
+
     def test_marimo_smc_progress(self):
         backend = MarimoProgressBackend(
             step_name="Stage", n_bars=1, total=1.0, combined=False, full_stats=False
@@ -139,3 +162,39 @@ def test_marimo_smc_progress(self):
             old = beta
 
         assert backend._task_state[0]["completed"] == 1.0
+
+    def test_nutpie_elapsed_freezes_after_completion(self):
+        """Completed nutpie chains must not show drifting elapsed in marimo."""
+        with patch("pymc.progress_bar.progress.in_marimo_notebook", return_value=True):
+            manager = NutpieProgressBarManager(chains=2, draws=100, progressbar=True)
+        assert isinstance(manager._backend, MarimoProgressBackend)
+        total = 1100
+
+        backend = manager._backend
+        backend._initialize_tasks()
+
+        def cp(finished, runtime_ms, started=True):
+            return SimpleNamespace(
+                finished_draws=finished,
+                total_draws=total,
+                runtime_ms=runtime_ms,
+                started=started,
+                divergent_draws=[],
+                step_size=0.5,
+                latest_num_steps=7,
+            )
+
+        # Chain 0 finishes, chain 1 halfway
+        manager.update([cp(total, runtime_ms=5000), cp(500, runtime_ms=2500)])
+
+        html_at_finish = backend._render_task_row(0, backend._task_state[0], [])
+
+        # Let wall-clock advance so unfrozen elapsed would visibly drift
+        sleep(0.3)
+
+        # More callbacks arrive while chain 1 is still running
+        manager.update([cp(total, runtime_ms=5000), cp(800, runtime_ms=4000)])
+
+        # Chain 0's row must be identical — elapsed and speed frozen
+        html_after = backend._render_task_row(0, backend._task_state[0], [])
+        assert html_at_finish == html_after