[data] Improve execution progress rendering

doc/source/data/monitoring-your-workload.rst

@@ -40,6 +40,9 @@ For operator names longer than a threshold of 100 characters, Ray Data truncates
 * To turn off this behavior and show the full operator name, set `DataContext.get_current().enable_progress_bar_name_truncation = False`.
 * To change the threshold of truncating the name, update the constant `ray.data._internal.progress_bar.ProgressBar.MAX_NAME_LENGTH = 42`.
 
+.. tip::
+   There is a new experimental console UI to show progress bars. Set `DataContext.get_current().enable_rich_progress_bars = True` or set the `RAY_DATA_ENABLE_RICH_PROGRESS_BARS=1` environment variable to enable.
+
 .. _ray-data-dashboard:
 
 Ray Data dashboard

python/ray/data/_internal/execution/operators/base_physical_operator.py

@@ -1,4 +1,4 @@
-import abc
+from abc import ABC, abstractmethod
 from typing import List, Optional
 
 from ray.data._internal.execution.interfaces import (
@@ -8,13 +8,14 @@
     TaskContext,
 )
 from ray.data._internal.execution.interfaces.physical_operator import _create_sub_pb
+from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
 from ray.data._internal.logical.interfaces import LogicalOperator
 from ray.data._internal.stats import StatsDict
 from ray.data.context import DataContext
 
 
-class InternalQueueOperatorMixin(PhysicalOperator, abc.ABC):
-    @abc.abstractmethod
+class InternalQueueOperatorMixin(PhysicalOperator, ABC):
+    @abstractmethod
     def internal_queue_size(self) -> int:
         """Returns Operator's internal queue size"""
         ...
@@ -47,7 +48,9 @@ def input_dependency(self) -> PhysicalOperator:
         return self.input_dependencies[0]
 
 
-class AllToAllOperator(InternalQueueOperatorMixin, PhysicalOperator):
+class AllToAllOperator(
+    InternalQueueOperatorMixin, SubProgressBarMixin, PhysicalOperator
+):
     """A blocking operator that executes once its inputs are complete.
 
     This operator implements distributed sort / shuffle operations, etc.
@@ -169,6 +172,15 @@ def close_sub_progress_bars(self):
             for sub_bar in self._sub_progress_bar_dict.values():
                 sub_bar.close()
 
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        return self._sub_progress_bar_names
+
+    def set_sub_progress_bar(self, name, pg):
+        # not type-checking due to circular imports
+        if self._sub_progress_bar_dict is None:
+            self._sub_progress_bar_dict = {}
+        self._sub_progress_bar_dict[name] = pg
+
     def supports_fusion(self):
         return True
 

python/ray/data/_internal/execution/operators/hash_shuffle.py

@@ -45,6 +45,7 @@
     _create_sub_pb,
     estimate_total_num_of_blocks,
 )
+from ray.data._internal.execution.operators.sub_progress import SubProgressBarMixin
 from ray.data._internal.logical.interfaces import LogicalOperator
 from ray.data._internal.stats import OpRuntimeMetrics
 from ray.data._internal.table_block import TableBlockAccessor
@@ -363,7 +364,7 @@ def combine(one: "_PartitionStats", other: "_PartitionStats") -> "_PartitionStat
         )
 
 
-class HashShuffleProgressBarMixin(abc.ABC):
+class HashShuffleProgressBarMixin(SubProgressBarMixin):
     @property
     @abc.abstractmethod
     def shuffle_name(self) -> str:
@@ -374,26 +375,29 @@ def shuffle_name(self) -> str:
     def reduce_name(self) -> str:
         ...
 
+    def _validate_sub_progress_bar_names(self):
+        assert self.shuffle_name is not None, "shuffle_name should not be None"
+        assert self.reduce_name is not None, "reduce_name should not be None"
+
     def initialize_sub_progress_bars(self, position: int) -> int:
-        """Display all sub progres bars in the termainl, and return the number of bars."""
+        """Display all sub progress bars in the termainl, and return the number of bars."""
+        self._validate_sub_progress_bar_names()
 
         # shuffle
         progress_bars_created = 0
         self.shuffle_bar = None
-        if self.shuffle_name is not None:
-            self.shuffle_bar, position = _create_sub_pb(
-                self.shuffle_name, self.num_output_rows_total(), position
-            )
-            progress_bars_created += 1
+        self.shuffle_bar, position = _create_sub_pb(
+            self.shuffle_name, self.num_output_rows_total(), position
+        )
+        progress_bars_created += 1
         self.shuffle_metrics = OpRuntimeMetrics(self)
 
         # reduce
         self.reduce_bar = None
-        if self.reduce_name is not None:
-            self.reduce_bar, position = _create_sub_pb(
-                self.reduce_name, self.num_output_rows_total(), position
-            )
-            progress_bars_created += 1
+        self.reduce_bar, position = _create_sub_pb(
+            self.reduce_name, self.num_output_rows_total(), position
+        )
+        progress_bars_created += 1
         self.reduce_metrics = OpRuntimeMetrics(self)
 
         return progress_bars_created
@@ -403,6 +407,27 @@ def close_sub_progress_bars(self):
         self.shuffle_bar.close()
         self.reduce_bar.close()
 
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        self._validate_sub_progress_bar_names()
+
+        # shuffle
+        self.shuffle_bar = None
+        self.shuffle_metrics = OpRuntimeMetrics(self)
+
+        # reduce
+        self.reduce_bar = None
+        self.reduce_metrics = OpRuntimeMetrics(self)
+
+        return [self.shuffle_name, self.reduce_name]
+
+    def set_sub_progress_bar(self, name, pg):
+        # No type-hints due to circular imports. `name` should be a `str`
+        # and `pg` should be a `SubProgressBar`
+        if self.shuffle_name is not None and self.shuffle_name == name:
+            self.shuffle_bar = pg
+        elif self.reduce_name is not None and self.reduce_name == name:
+            self.reduce_bar = pg
+
 
 def _derive_max_shuffle_aggregators(
     total_cluster_resources: ExecutionResources,
@@ -705,7 +730,7 @@ def _on_partitioning_done(cur_shuffle_task_idx: int):
                 self.shuffle_metrics.on_task_finished(cur_shuffle_task_idx, None)
 
                 # Update Shuffle progress bar
-                self.shuffle_bar.update(i=input_block_metadata.num_rows)
+                self.shuffle_bar.update(increment=input_block_metadata.num_rows or 0)
 
             # TODO update metrics
             task = self._shuffling_tasks[input_index][
@@ -823,7 +848,7 @@ def _on_bundle_ready(partition_id: int, bundle: RefBundle):
 
             # Update Finalize progress bar
             self.reduce_bar.update(
-                i=bundle.num_rows(), total=self.num_output_rows_total()
+                increment=bundle.num_rows() or 0, total=self.num_output_rows_total()
             )
 
         def _on_aggregation_done(partition_id: int, exc: Optional[Exception]):

python/ray/data/_internal/execution/operators/sub_progress.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+
+class SubProgressBarMixin(ABC):
+    """Abstract class for operators that support sub-progress bars"""
+
+    @abstractmethod
+    def get_sub_progress_bar_names(self) -> Optional[List[str]]:
+        """
+        Returns list of sub-progress bar names
+
+        This is used to create the sub-progress bars in the progress manager.
+        Note that sub-progress bars will be created in the order returned by
+        this method.
+        """
+        ...
+
+    @abstractmethod
+    def set_sub_progress_bar(self, name, pg):
+        """
+        Sets sub-progress bars
+
+        name: name of sub-progress bar
+        pg: SubProgressBar instance (progress_manager.py)
+        """
+        # Skipping type-checking for circular imports
+        ...