[data] Improve execution progress rendering

[kyuds]

Why are these changes needed?

Link to video recording

Currently, by default, the original tqdm based progress is used. To enable rich progress reporting as shown in the screenshot, set:

ray.data.DataContext.get_current().enable_rich_progress_bars = True

or set the envvar:

export RAY_DATA_ENABLE_RICH_PROGRESS_BARS=1

Related issue number

Fixes #52505

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

---

Note

Replaces legacy per-operator/global progress bars with a Rich-based progress manager that tracks global/operator progress and resources, refactors topology/progress APIs, and updates tests.

• Execution Progress (Rich-based):
  • Introduces progress_manager.py with RichExecutionProgressManager for global/operator progress, rates, elapsed/remaining time, and live resource usage.
  • Streaming executor integrates manager (start/refresh/close, finishing messages), updates on row/output and resources, and periodic refresh via PROGRESS_MANAGER_UPDATE_INTERVAL.
• Operator/State Refactor:
  • OpState gains OpDisplayMetrics, progress_manager_uuid, output_row_count, and update_display_metrics; removes legacy progress bar handling and summary methods.
  • _debug_dump_topology now logs op_display_metrics.display_str().
  • Minor: add TODOs on sub-progress-bar helpers in AllToAllOperator and HashShuffleProgressBarMixin.
• Topology API:
  • build_streaming_topology(...) now returns only Topology (no progress bar count); all call sites and tests updated.
• Iterator/Reporting:
  • _ClosingIterator updates total progress via manager.
  • Resource reporting moved to progress manager.
• Tests:
  • Adjust unit tests to new topology return type and removed progress bar expectations.

^{Written by Cursor Bugbot for commit 935f3c3. This will update automatically on new commits. Configure here.}

[gemini-code-assist]

Code Review

This pull request introduces a significant improvement to the execution progress rendering by replacing the custom ProgressBar with the rich library. This provides a much more modern and informative UI. The refactoring is extensive, touching the streaming executor, its state management, and operators.
The changes are generally well-implemented, but I've found a couple of critical issues that need to be addressed. One is a bug in how progress is updated, which would lead to incorrect progress display. The other is related to the WIP nature of the PR, where a missing dependency check for rich would cause a crash.
Additionally, there are several TODO comments, especially TODO (kyuds): remove, which should be cleaned up before this PR is merged.
Overall, this is a great enhancement, and with the suggested fixes, it will be a solid contribution.

[kyuds]

cc @richardliaw for initial feedback!

[goutamvenkat-anyscale]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new rich-based execution progress manager, which is a significant improvement over the legacy progress bars. The refactoring across the streaming executor and operator state is well-executed. I've identified a few issues: a logic bug in progress metric calculation that could lead to incorrect progress display, a redundant stats generation call, and a minor type safety issue in the new OpDisplayMetrics dataclass. Overall, this is a great enhancement to Ray Data's user experience.

[goutamvenkat-anyscale]

Thanks for the change! It's much needed. Also love the decoupling of the progress bar from the streaming executor.

[kyuds]

I checked out the buildkite failures and they fall largely into two types:

For failure type 1:
As we consolidate progress reporting, we no longer support sub-progress bars that were initialized in initialize_sub_progress_bars. The errors for argument of type "NoneType" is non-iterable stems from that we try to get the sub-progress bars that we no longer use, hence we no longer initialize.

These changes were introduced in #34997.

I see two solutions for this:

1. We remove the sub progress bars for good --> that is inline with the "simplification" of progress reporting
2. We "turn off" sub progress bars for the context of this PR, and then we improve the UI to introduce sub-progress bars to the current layout (reserved for a future pr)

I am open to either method

For failure type 2:
Very similar to failure type 1, but now the error message is shuffle_metrics attribute doesn't exist. This happens because in initialize_sub_progress_bars, we also initialize OpRuntimeMetrics to I guess track the sub-progress of things. The only problem is, while OpRuntimeMetrics is a dataclass, it creates internal bundle queues that seems to store objects.

A potential solution I see:
Simply just only initialize OpRuntimeMetrics in a function like initialize_metrics, and check if our operator instance is of type HashShuffleProgressBarMixin, and initialize the metrics as appropriate.

I am super new to the codebase, and I am more than welcome to any feedback/suggestions/etc!!

cc @goutamvenkat-anyscale , @richardliaw

[kyuds]

Just checking back in to see if theres anything I could do!

[goutamvenkat-anyscale]

I checked out the buildkite failures and they fall largely into two types:

For failure type 1: As we consolidate progress reporting, we no longer support sub-progress bars that were initialized in initialize_sub_progress_bars. The errors for argument of type "NoneType" is non-iterable stems from that we try to get the sub-progress bars that we no longer use, hence we no longer initialize.

These changes were introduced in #34997.

I see two solutions for this:

1. We remove the sub progress bars for good --> that is inline with the "simplification" of progress reporting
2. We "turn off" sub progress bars for the context of this PR, and then we improve the UI to introduce sub-progress bars to the current layout (reserved for a future pr)

I am open to either method

For failure type 2: Very similar to failure type 1, but now the error message is shuffle_metrics attribute doesn't exist. This happens because in initialize_sub_progress_bars, we also initialize OpRuntimeMetrics to I guess track the sub-progress of things. The only problem is, while OpRuntimeMetrics is a dataclass, it creates internal bundle queues that seems to store objects.

A potential solution I see: Simply just only initialize OpRuntimeMetrics in a function like initialize_metrics, and check if our operator instance is of type HashShuffleProgressBarMixin, and initialize the metrics as appropriate.

I am super new to the codebase, and I am more than welcome to any feedback/suggestions/etc!!

cc @goutamvenkat-anyscale , @richardliaw

@kyuds Sorry for the delay. We use sub progress bars for all shuffle operations, cause there are multiple stages (like map and reduce). So simply turning off the bars is not a viable option imo. Can you estimate how much work it is to incorporate the management of sub progress bars inside your new progress bar manager?

I believe it should be these operators (but please double check if there are more):

• AllToAllOperator.initialize_sub_progress_bars()
• HashShuffleProgressBarMixin.initialize_sub_progress_bars()
• OpState.initialize_progress_bars()

Or another option is we break it up into 2 PRs, in this one, we let the sub progress bars use tqdm/the current PB manager and then in the subsequent PR, you can make the set of changes needed to refactor subprogess bars to use rich and then deprecate the old manager.

As for the metics initialization I'm fine with decoupling that as well, as you suggested.

[kyuds]

Thank you for the feedback @goutamvenkat-anyscale !

I think its going to take me one to two days to make the subprogress bars work (however, as this was not part of the original spec, I would really appreciate getting feedback on how the sub-progress bars should actually look like. I'll do an initial design and I think we can exchange feedback and iterate from there).

For decoupling the metrics, I think that can come later actually, otherwise the PR will become extremely bloated.

The only problem for me right now is I recently had a medical operation so I can't even put on my glasses... (Im literally squinting hard at my screen trying to type this reply). I think I'll be able to implement the fixes starting early next-week and finish by end of mid week. Hope you understand the delay 🙏 .

[gvspraveen]

@kyuds Thanks for the contribution. No worries. Hope you feel better soon! Take care!

[kyuds]

@iamjustinhsu pushed the fix for the rate string! Thank you so much!

[iamjustinhsu]

nice work!! :)

[alexeykudinkin]

Why is this reinitializing the metrics and resetting the bar?

[kyuds]

this is called once upon initializing progress manager, not called at all when using the tqdm progress bar implementation.

btw, when initializing the progress manager, the initialize_sub_progress_bars method is not called, so we need an alternative method of initializing the metrics, which is done in that function. Otherwise the program will crash entirely because metrics are not initialized at all.

However, I do understand the naming of the function may cause confusion. Would you have some suggestions for an alternative name?

[kyuds]

some ideas that come into mind: initialize_sub_progress_bar_infos, initialize_sub_progress_attrs
I'll submit a follow-up PR for this if any is ok with you. Any suggestions are, as always, very welcome

[alexeykudinkin]

@kyuds thank you for your great contribution!

As a follow-up would you mind actually doing the following:

• Abstracting both progress-bar implementations away completely from the StreamingExecutor behind ProgressManager interface
• Maintain 2 implementations (Rich-basd and tqdm) that could be extending ProgressManager (or being injected into)

[alexeykudinkin]

There's no need to set shuffle_bar to none and immediately overwrite it

[alexeykudinkin]

Add typing

[alexeykudinkin]

This should be

if self._use_rich_progress():
   ... 
elif self._global_info:
   ...