Skip to content

Commit

Permalink
Fix Rich progress bar flickering by limiting refreshes (#571)
Browse files Browse the repository at this point in the history
  • Loading branch information
tomwhite authored Sep 9, 2024
1 parent 7b32304 commit 88ae813
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion cubed/diagnostics/rich.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
import sys
import time
from contextlib import contextmanager

from rich.console import RenderableType
Expand All @@ -18,6 +19,8 @@
from cubed.runtime.pipeline import visit_nodes
from cubed.runtime.types import Callback

REFRESH_PER_SECOND = 10


class RichProgressBar(Callback):
"""Rich progress bar for a computation."""
Expand Down Expand Up @@ -50,6 +53,7 @@ def on_compute_start(self, event):
self.logger_aware_progress = logger_aware_progress
self.progress = progress
self.progress_tasks = progress_tasks
self.last_updated = time.time()

def on_compute_end(self, event):
self.logger_aware_progress.__exit__(None, None, None)
Expand All @@ -58,8 +62,11 @@ def on_operation_start(self, event):
self.progress.start_task(self.progress_tasks[event.name])

def on_task_end(self, event):
now = time.time()
refresh = now - self.last_updated > (1.0 / REFRESH_PER_SECOND)
self.last_updated = now
self.progress.update(
self.progress_tasks[event.name], advance=event.num_tasks, refresh=True
self.progress_tasks[event.name], advance=event.num_tasks, refresh=refresh
)


Expand Down

0 comments on commit 88ae813

Please sign in to comment.