import dataclasses import logging import math import sys import time import typing from typing import Dict, List, Optional, Tuple from rich.console import Console from rich.live import Live from rich.progress import ( BarColumn, Progress, SpinnerColumn, TaskID, TextColumn, TimeElapsedColumn, ) from rich.table import Column, Table from rich.text import Text from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin from ray.data._internal.execution.streaming_executor_state import ( format_op_state_summary, ) from ray.data._internal.progress.base_progress import ( BaseExecutionProgressManager, BaseProgressBar, NoopSubProgressBar, ) from ray.data._internal.progress.utils import truncate_operator_name if typing.TYPE_CHECKING: from ray.data._internal.execution.resource_manager import ResourceManager from ray.data._internal.execution.streaming_executor_state import OpState, Topology logger = logging.getLogger(__name__) _TREE_BRANCH = " ├─" _TREE_VERTICAL = "│" _TREE_VERTICAL_SUB_PROGRESS = " │ -" _TREE_VERTICAL_INDENT = f" {_TREE_VERTICAL} " _TOTAL_PROGRESS_TOTAL = 1.0 _RESOURCE_REPORT_HEADER = f" {_TREE_VERTICAL} Active/total resources: " class RichSubProgressBar(BaseProgressBar): """Thin wrapper to provide identical interface to the ProgressBar. Updates RichExecutionProgressManager internally. """ def __init__( self, name: str, total: Optional[int] = None, progress: Progress = None, tid: TaskID = None, max_name_length: int = 100, ): """ Initialize sub-progress bar Args: name: name of sub-progress bar total: total number of output rows. None for unknown. progress: rich.Progress instance for the corresponding sub-progress bar. tid: rich.TaskId for the corresponding sub-progress bar task. max_name_length: maximum operator name length. """ self._total = total self._completed = 0 self._start_time = None self._enabled = True self._progress = progress self._tid = tid self._max_name_length = max_name_length self._desc = truncate_operator_name(name, self._max_name_length) def set_description(self, name: str) -> None: self._desc = truncate_operator_name(name, self._max_name_length) if self._enabled: self._progress.update(self._tid, description=self._desc) def get_description(self) -> str: return self._desc def _update(self, completed: int, total: Optional[int] = None) -> None: assert self._enabled if self._start_time is None: self._start_time = time.time() metrics = _get_progress_metrics(self._start_time, completed, total) self._progress.update( self._tid, completed=metrics.completed, total=metrics.total, rate_str=metrics.rate_str, count_str=metrics.count_str, ) def update(self, increment: int = 0, total: Optional[int] = None) -> None: if self._enabled and increment != 0: if total is not None: self._total = total self._completed += increment self._update(self._completed, self._total) def complete(self) -> None: if self._enabled: self._update(self._completed, self._completed) def __getstate__(self): return {"max_name_length": self._max_name_length} def __setstate__(self, state): self._enabled = False # Progress bar is disabled on remote nodes. self._max_name_length = state.get("max_name_length", 100) class RichExecutionProgressManager(BaseExecutionProgressManager): """Execution progress display using rich.""" def __init__( self, dataset_id: str, topology: "Topology", show_op_progress: bool, verbose_progress: bool, ): self._dataset_id = dataset_id self._sub_progress_bars: List[BaseProgressBar] = [] self._show_op_progress = show_op_progress self._verbose_progress = verbose_progress self._start_time: Optional[float] = None # rich self._console = Console(file=sys.stderr) self._total = self._make_progress_bar(" ", "•", 15) self._current_rows = 0 self._total_resources = Text( f"{_RESOURCE_REPORT_HEADER}Initializing...", no_wrap=True ) self._op_display: Dict[ "OpState", Tuple[Optional[TaskID], Optional[Progress], Optional[Text]] ] = {} self._layout_table = Table.grid(padding=(0, 1, 0, 0), expand=True) self._layout_table.add_row(self._total) self._layout_table.add_row(self._total_resources) self._setup_operator_progress(topology) # empty new line to prevent "packed" feeling self._layout_table.add_row(Text()) # rich.Live is the auto-refreshing rich component display. # refreshing/closing is all done through rich.Live self._live = Live( self._layout_table, console=self._console, refresh_per_second=2, vertical_overflow="visible", ) self._total_task_id = self._total.add_task( f"Dataset {self._dataset_id} running:", total=_TOTAL_PROGRESS_TOTAL, rate_str="? rows/s", count_str="0/?", ) def _setup_operator_progress(self, topology: "Topology"): rows = [] for state in topology.values(): op = state.op if isinstance(op, InputDataBuffer): continue contains_sub_progress_bars = isinstance(op, SubProgressBarMixin) sub_progress_bar_enabled = self._show_op_progress and ( contains_sub_progress_bars or self._verbose_progress ) if sub_progress_bar_enabled: progress = self._make_progress_bar(_TREE_BRANCH, " ", 10) stats = Text(f"{_TREE_VERTICAL_INDENT}Initializing...", no_wrap=True) total = state.op.num_output_rows_total() name = truncate_operator_name(state.op.name, self.MAX_NAME_LENGTH) tid = progress.add_task( name, total=total if total is not None else 1, start=True, rate_str="? rows/s", count_str="0/?", ) rows.append(progress) rows.append(stats) self._op_display[state] = (tid, progress, stats) if not contains_sub_progress_bars: continue sub_progress_bar_names = op.get_sub_progress_bar_names() if sub_progress_bar_names is None: continue for name in sub_progress_bar_names: if sub_progress_bar_enabled: progress = self._make_progress_bar( _TREE_VERTICAL_SUB_PROGRESS, "", 10 ) total = state.op.num_output_rows_total() tid = progress.add_task( name, total=total if total is not None else 1, start=True, rate_str="? rows/s", count_str="0/?", ) rows.append(progress) pg = RichSubProgressBar( name=name, total=total, progress=progress, tid=tid, max_name_length=self.MAX_NAME_LENGTH, ) else: pg = NoopSubProgressBar( name=name, max_name_length=self.MAX_NAME_LENGTH ) op.set_sub_progress_bar(name, pg) self._sub_progress_bars.append(pg) if rows: self._layout_table.add_row(Text(f" {_TREE_VERTICAL}", no_wrap=True)) for row in rows: self._layout_table.add_row(row) def _make_progress_bar(self, indent_str: str, spinner_finish: str, bar_width: int): return Progress( TextColumn(indent_str, table_column=Column(no_wrap=True)), SpinnerColumn(finished_text=spinner_finish), TextColumn( "{task.description} {task.percentage:>3.0f}%", table_column=Column(no_wrap=True), ), BarColumn(bar_width=bar_width), TextColumn("{task.fields[count_str]}", table_column=Column(no_wrap=True)), TextColumn("["), TimeElapsedColumn(), TextColumn(","), TextColumn("{task.fields[rate_str]}", table_column=Column(no_wrap=True)), TextColumn("]"), console=self._console, transient=False, expand=False, ) # Management def start(self): if not self._live.is_started: self._live.start() def refresh(self): if self._live.is_started: self._live.refresh() def close_with_finishing_description(self, desc: str, success: bool): if self._live.is_started: kwargs = {} if success: # set everything to completed kwargs["completed"] = 1.0 kwargs["total"] = 1.0 for pg in self._sub_progress_bars: if isinstance(pg, RichSubProgressBar): pg.complete() if self._start_time is None: self._start_time = time.time() for tid, progress, _ in self._op_display.values(): completed = progress.tasks[tid].completed or 0 metrics = _get_progress_metrics( self._start_time, completed, completed ) _update_with_conditional_rate(progress, tid, metrics) self._total.update(self._total_task_id, description=desc, **kwargs) self.refresh() # need this sleep delay to ensure that changes are rendered to screen # before rich Live module is stopped. time.sleep(0.02) self._live.stop() # Total Progress def _can_update_total(self) -> bool: return ( self._total_task_id is not None and self._total_task_id in self._total.task_ids ) def update_total_progress(self, new_rows: int, total_rows: Optional[int]): if not self._can_update_total(): return if self._live.is_started: if self._start_time is None: self._start_time = time.time() if new_rows is not None: self._current_rows += new_rows metrics = _get_progress_metrics( self._start_time, self._current_rows, total_rows ) _update_with_conditional_rate(self._total, self._total_task_id, metrics) def update_total_resource_status(self, resource_status: str): if not self._can_update_total(): return if self._live.is_started: self._total_resources.plain = _RESOURCE_REPORT_HEADER + resource_status def _can_update_operator(self, op_state: "OpState") -> bool: if op_state not in self._op_display: return False tid, progress, stats = self._op_display[op_state] if tid is None or not progress or not stats or tid not in progress.task_ids: return False return True def update_operator_progress( self, op_state: "OpState", resource_manager: "ResourceManager" ): if not self._can_update_operator(op_state): return if self._start_time is None: self._start_time = time.time() tid, progress, stats = self._op_display[op_state] # progress current_rows = op_state.op.metrics.row_outputs_taken total_rows = op_state.op.num_output_rows_total() metrics = _get_progress_metrics(self._start_time, current_rows, total_rows) _update_with_conditional_rate(progress, tid, metrics) # stats stats_str = format_op_state_summary(op_state, resource_manager) stats.plain = f"{_TREE_VERTICAL_INDENT}{stats_str}" # utilities def _format_k(val: int) -> str: if val >= 1000: fval = val / 1000.0 fval_str = f"{int(fval)}" if fval.is_integer() else f"{fval:.2f}" return fval_str + "k" return str(val) def _format_row_count(completed: int, total: Optional[int]) -> str: """Formats row counts with k units.""" cstr = _format_k(completed) if total is None or math.isinf(total): tstr = "?k" if cstr.endswith("k") else "?" else: tstr = _format_k(total) return f"{cstr}/{tstr}" @dataclasses.dataclass class _ProgressMetrics: completed: int total: int rate_str: str count_str: str def _get_progress_metrics( start_time: float, completed_rows: int, total_rows: Optional[int] ) -> _ProgressMetrics: """ Args: start_time: time when progress tracking started completed_rows: cumulative rows outputted total_rows: total rows expected (can be unknown) Returns: _ProgressMetrics instance containing the calculated data. """ # Note, when total is unknown, we default the progress bar to 0. # Will properly have estimates for rate and count strings though. total = 1 if total_rows is None or total_rows < 1 else total_rows completed = 0 if total_rows is None else completed_rows if total_rows is None: rate_str = "? row/s" else: elapsed = time.time() - start_time rate_val = completed_rows / elapsed if elapsed > 1 else 0 rate_unit = "row/s" if rate_val >= 1000: rate_val /= 1000 rate_unit = "k row/s" rate_str = f"{rate_val:.2f} {rate_unit}" count_str = _format_row_count(completed_rows, total_rows) return _ProgressMetrics( completed=completed, total=total, rate_str=rate_str, count_str=count_str ) def _update_with_conditional_rate( progress: Progress, tid: TaskID, metrics: _ProgressMetrics ): # not doing type checking because rich is imported conditionally. # progress: rich.Progress # tid: rich.TaskId # metrics: _ProgressMetrics task = progress.tasks[tid] kwargs = { "completed": metrics.completed, "total": metrics.total, "count_str": metrics.count_str, } if task.completed != metrics.completed: # update rate string only if there are new rows. # this allows updates to other metric data while # preserving the right rate notation. kwargs["rate_str"] = metrics.rate_str progress.update(tid, **kwargs)