[Data] Fix the config to disable progress bar (#44342)

This is a followup of #43360 to fix the behavior of disabling progress bar. After this PR, users only need to set `ray.data.DataContext.get_current().enable_progress_bars = False` to disable all progress bar (main + verbose bar per operator). Tested locally and verified config is working. Signed-off-by: Cheng Su <[email protected]>
ray-project · Apr 17, 2024 · a130313 · a130313
1 parent d8b35ed
commit a130313
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/python/ray/data/_internal/execution/streaming_executor_state.py b/python/ray/data/_internal/execution/streaming_executor_state.py
@@ -34,6 +34,7 @@
 from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
 from ray.data._internal.execution.resource_manager import ResourceManager
 from ray.data._internal.progress_bar import ProgressBar
+from ray.data.context import DataContext
 
 logger = logging.getLogger(__name__)
 
@@ -197,7 +198,9 @@ def initialize_progress_bars(self, index: int, verbose_progress: bool) -> int:
  """
  is_all_to_all = isinstance(self.op, AllToAllOperator)
  # Only show 1:1 ops when in verbose progress mode.
- enabled = verbose_progress or is_all_to_all
+ enabled = is_all_to_all or (
+ DataContext.get_current().enable_progress_bars and verbose_progress
+ )
  self.progress_bar = ProgressBar(
  "- " + self.op.name,
  self.op.num_outputs_total(),