Add progress bars to hash operators #53175

iamjustinhsu · 2025-05-20T17:19:02Z

Why are these changes needed?

We want to make transparent progress bars for all operators. Currently, join operators (or any operator that uses HashShufflingBaseOperator) will not show a progress bar, so it looks a lil funky.

After these changes

Other changes

In OpruntimeMetrics, added num_row_inputs_received (a counterpart to num_inputs_received)
Uplifted some estimated row output computation logic from the MapOperator to PhysicalOperator so it can also be used by HashShufflingBaseOperator

Related issue number

https://anyscale1.atlassian.net/browse/DATA-574 and https://anyscale1.atlassian.net/browse/DATA-925

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 :(

…u/add-pb-to-hash-operators

alexeykudinkin · 2025-05-27T22:00:02Z

python/ray/data/_internal/execution/interfaces/physical_operator.py

+        self._sub_progress_bar_names = sub_progress_bar_names
+        self._sub_progress_bar_dict: Dict[str, ProgressBar] = None


Do we really need names separate form PBs themselves?

If you initialize a ProgressBar, it will display in the terminal. we don't want this behavior because it will appear as if the operation is running, but it's not, until we initialize all sub progress bars

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

alexeykudinkin · 2025-05-27T22:08:31Z

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

+            self._metrics.on_task_output_generated(task_index=task_index, output=bundle)
            self._output_queue.append(bundle)
+            shuffle_reduce_bar.update(
+                i=bundle.num_rows(), total=self.num_output_rows_total()
+            )
            self._metrics.on_output_queued(bundle)


nit: Please move metrics callbacks together

sure no strong opinion on this instance, but following this example, it seems like self._metrics.on_task_output_generated(task_index=task_index, output=bundle) is on the start of the function and self._metrics.on_output_queued(bundle) is at the end of the function

python/ray/data/_internal/execution/interfaces/physical_operator.py

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

alexeykudinkin · 2025-05-28T22:07:05Z

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

+            self._metrics.on_task_submitted(
+                self._next_shuffle_tasks_idx + partition_id,
+                RefBundle([], owns_blocks=True),
+            )
+            shuffle_reduce_bar.update(i=0, total=self.num_output_rows_total())


So this will be using single estimate for both shuffle and finalization tasks. Mixing them up doesn't really make a lot of sense.

What i'm thinking is we should explore is using separate metrics for shuffling and finalization stages (which will be kind of split the operator into 2)

…u/add-pb-to-hash-operators

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

This reverts commit 793c03f.

This reverts commit 74f00c2.

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin · 2025-06-26T01:18:27Z

python/ray/data/_internal/execution/interfaces/physical_operator.py


 class ReportsExtraResourceUsage(abc.ABC):
    @abc.abstractmethod
    def extra_resource_usage(self: PhysicalOperator) -> ExecutionResources:
        """Returns resources used by this operator beyond standard accounting."""
        ...
+
+
+class ContainsSubProgressBars(PhysicalOperator):


Suggested change

class ContainsSubProgressBars(PhysicalOperator):

class WithSubProgressBarMixin:

Why do we need to inherit from PO?

i was doing it for type checking, i can change it back

alexeykudinkin · 2025-06-26T01:20:57Z

python/ray/data/_internal/execution/interfaces/physical_operator.py

+        super().__init__(*args, **kwargs)
+        self._sub_progress_bar_names: Optional[List[str]] = sub_progress_bar_names
+        self._sub_progress_bar_dict: Optional[Dict[str, ProgressBar]] = None
+        self._metric_dict: Dict[str, OpRuntimeMetrics] = {}


I think this should live outside of this one -- let's start by keeping it in the Shuffling Operator base but if there's opportunity to abstract away we'd see it and do it later

alexeykudinkin · 2025-06-26T01:26:25Z

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

@@ -334,7 +347,7 @@ def combine(one: "_PartitionStats", other: "_PartitionStats") -> "_PartitionStat
        )


-class HashShufflingOperatorBase(PhysicalOperator):
+class HashShufflingOperatorBase(ContainsSubProgressBars, PhysicalOperator):


Suggested change

class HashShufflingOperatorBase(ContainsSubProgressBars, PhysicalOperator):

class HashShufflingOperatorBase(PhysicalOperator, ContainsSubProgressBars):

alexeykudinkin · 2025-06-26T01:27:26Z

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

+            shuffle_progress_bar_name = "Hash Shuffle Map"
+        if finalize_progress_bar_name is None:
+            finalize_progress_bar_name = "Hash Shuffle Reduce"


Suggested change

shuffle_progress_bar_name = "Hash Shuffle Map"

if finalize_progress_bar_name is None:

finalize_progress_bar_name = "Hash Shuffle Reduce"

shuffle_progress_bar_name = "Shuffle"

if finalize_progress_bar_name is None:

finalize_progress_bar_name = "Reduce"

Let's just simplify it

alexeykudinkin · 2025-06-26T01:28:38Z

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

@@ -451,90 +476,129 @@ def start(self, options: ExecutionOptions) -> None:
        self._aggregator_pool.start()

    def _add_input_inner(self, input_bundle: RefBundle, input_index: int) -> None:
+
+        shuffle_metrics = self.get_metrics(0)


Following up with my comment above -- let's re-home these multi-metrics set up in this class (being _shuffle_stage_metrics and _reduce_stage_metrics)

alexeykudinkin · 2025-06-26T01:39:41Z

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

-                cur_shuffle_task_idx % self._num_partitions
-                if not input_key_column_names
-                else None
+        def _on_partitioning_done(cur_shuffle_task_idx: int):


I see you're pulling up this method and unfortunately now GH just shows this whole file as pretty much rewritten because of shifting lines and some other changes laid on top.

It's totally cool to refactor if you want to make things better but please avoid mixing up refactorings like that and critical changes in semantic. It's making it practically impossible for me to effectively review this change and assert correctness of what we're doing.

Let's extract the refactoring part into PR stacked on top so these could be reviewed separately (and much faster)

got it, ill try to move it back to what it was before

didn't intend to make it harder to review

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch 5 times, most recently from 02ff932 to 16f616d Compare May 21, 2025 16:51

Add progress bars to hash operators

62abdd3

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch 2 times, most recently from 587538c to 5275fa6 Compare May 21, 2025 23:44

simplify attributes -> names

00a6ead

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch from 5275fa6 to 00a6ead Compare May 21, 2025 23:45

hainesmichaelc added the community-backlog label May 22, 2025

Merge branch 'ray-project:master' into jhsu/add-pb-to-hash-operators

a6e4219

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch from 682b0ba to a6e4219 Compare May 22, 2025 16:40

hainesmichaelc removed the community-backlog label May 22, 2025

Merge branch 'master' of https://github.com/iamjustinhsu/ray into jhs…

bc84166

…u/add-pb-to-hash-operators

iamjustinhsu marked this pull request as ready for review May 27, 2025 18:49

iamjustinhsu requested a review from a team as a code owner May 27, 2025 18:49

Merge branch 'ray-project:master' into jhsu/add-pb-to-hash-operators

733b1c0

alexeykudinkin reviewed May 28, 2025

View reviewed changes

iamjustinhsu and others added 5 commits May 28, 2025 15:43

Merge branch 'ray-project:master' into jhsu/add-pb-to-hash-operators

5c3c2e5

Merge branch 'master' of https://github.com/iamjustinhsu/ray into jhs…

0767a6a

…u/add-pb-to-hash-operators

Merge branch 'master' of https://github.com/iamjustinhsu/ray into jhs…

8db7362

…u/add-pb-to-hash-operators

add mixin for sub progress bars

78276aa

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

fix tests

c6e1abe

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch from 4427aa7 to c6e1abe Compare June 12, 2025 01:17

is this better

74f00c2

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch 2 times, most recently from 74f00c2 to 6879943 Compare June 12, 2025 04:14

iamjustinhsu requested review from a team and edoakes as code owners June 12, 2025 04:14

iamjustinhsu removed request for a team, pcmoritz, thomasdesr, jjyao, raulchen, richardliaw, edoakes, kevin85421, SongGuyang and aslonnie June 12, 2025 04:37

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch 2 times, most recently from 7d97483 to 7058ede Compare June 12, 2025 14:51

revert this commit to generatlize

793c03f

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch from 7058ede to 793c03f Compare June 12, 2025 14:52

iamjustinhsu added 2 commits June 12, 2025 16:59

Revert "revert this commit to generatlize"

6ae137b

This reverts commit 793c03f.

Revert "is this better"

9a3d0a9

This reverts commit 74f00c2.

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch from 924df8c to 9bf7929 Compare June 13, 2025 00:34

cleaner but makes more sense

1fe7991

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch 2 times, most recently from ed06ea8 to 2d3b83d Compare June 13, 2025 02:15

finalize task idx

fac4a70

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch 2 times, most recently from 455bcaf to 044bbf8 Compare June 13, 2025 02:49

final

251edfd

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu force-pushed the jhsu/add-pb-to-hash-operators branch from 044bbf8 to 251edfd Compare June 13, 2025 16:33

alexeykudinkin reviewed Jun 26, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add progress bars to hash operators #53175

Add progress bars to hash operators #53175

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

		self._sub_progress_bar_names = sub_progress_bar_names
		self._sub_progress_bar_dict: Dict[str, ProgressBar] = None

	class ContainsSubProgressBars(PhysicalOperator):
	class WithSubProgressBarMixin:

	class HashShufflingOperatorBase(ContainsSubProgressBars, PhysicalOperator):
	class HashShufflingOperatorBase(PhysicalOperator, ContainsSubProgressBars):

Add progress bars to hash operators #53175

Are you sure you want to change the base?

Add progress bars to hash operators #53175

Uh oh!

Conversation

Uh oh!

Why are these changes needed?

After these changes

Other changes

Related issue number

Checks

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!