[Data] Replace number of bundles with proper number of blocks by alexeykudinkin · Pull Request #58030 · ray-project/ray

alexeykudinkin · 2025-10-23T01:06:04Z

Description

Currently, we implicitly assume that RefBundle holds exactly 1 block. That's not a safe assumption, and this change is addressing that by explicitly referring to number of blocks instead

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

gemini-code-assist

Code Review

This pull request is a good step towards making block counting more accurate by explicitly using the number of blocks instead of assuming one block per bundle. The refactoring is applied across many files. However, I've found a few critical issues where the new counting logic is implemented incorrectly, which could lead to bugs in metrics, backpressure, and autoscaling. I've also included some minor suggestions for improving code style and efficiency.

python/ray/data/_internal/execution/bundle_queue/fifo_bundle_queue.py

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

.../ray/data/_internal/execution/backpressure_policy/downstream_capacity_backpressure_policy.py

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

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

bveeramani · 2025-10-24T22:00:01Z

python/ray/tests/test_runtime_env_working_dir.py

    @ray.remote
    def test_import():
        import file_module
+


Unrelated change for this and python/ray/tests/test_runtime_env_working_dir_3.py?

iamjustinhsu · 2025-10-24T22:21:52Z

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

        self._bundle_buffer_size = 0
+        self._bundle_buffer_size_bytes = 0


what's the difference between these 2 vars?

iamjustinhsu · 2025-10-24T22:22:42Z

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

-    def internal_queue_size(self) -> int:
-        return len(self._buffer)
+    def internal_queue_num_blocks(self) -> int:
+        return sum(len(b.block_refs) for b in self._buffer)


if b is type RefBundle, you can do len(b)

…bundle Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

… num bytes held inside internal queue; Revisited `InternalQueueOperatorMixin.internal_queue_size` to remove assumption that every bundle holds just 1 block Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing tests Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

…oject#58030) ## Description Currently, we implicitly assume that `RefBundle` holds exactly 1 block. That's not a safe assumption, and this change is addressing that by explicitly referring to number of blocks instead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

…oject#58030) ## Description Currently, we implicitly assume that `RefBundle` holds exactly 1 block. That's not a safe assumption, and this change is addressing that by explicitly referring to number of blocks instead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Aydin Abiar <aydin@anyscale.com>

…oject#58030) ## Description Currently, we implicitly assume that `RefBundle` holds exactly 1 block. That's not a safe assumption, and this change is addressing that by explicitly referring to number of blocks instead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: Future-Outlier <eric901201@gmail.com>

…oject#58030) ## Description Currently, we implicitly assume that `RefBundle` holds exactly 1 block. That's not a safe assumption, and this change is addressing that by explicitly referring to number of blocks instead ## Related issues > Link related issues: "Fixes ray-project#1234", "Closes ray-project#1234", or "Related to ray-project#1234". ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: Alexey Kudinkin <ak@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

alexeykudinkin requested review from a team as code owners October 23, 2025 01:06

gemini-code-assist bot reviewed Oct 23, 2025

View reviewed changes

This comment was marked as outdated.

Sign in to view

ray-gardener bot added the data Ray Data-related issues label Oct 23, 2025

alexeykudinkin added the go add ONLY when ready to merge, run all tests label Oct 23, 2025

alexeykudinkin force-pushed the ak/bndl-blk-fix branch from bb87078 to ea982b3 Compare October 24, 2025 06:48

bveeramani approved these changes Oct 24, 2025

View reviewed changes

bveeramani mentioned this pull request Oct 24, 2025

[Data] Revisiting OpResourceAllocator to make data flow explicit #57788

Merged

16 tasks

iamjustinhsu reviewed Oct 24, 2025

View reviewed changes

iamjustinhsu approved these changes Oct 24, 2025

View reviewed changes

alexeykudinkin enabled auto-merge (squash) October 27, 2025 05:05

alexeykudinkin force-pushed the ak/bndl-blk-fix branch from 5379f9c to e8ab2c8 Compare October 27, 2025 05:05

github-actions bot disabled auto-merge October 27, 2025 05:05

alexeykudinkin enabled auto-merge (squash) October 27, 2025 05:55

github-actions bot disabled auto-merge October 27, 2025 17:16

This comment was marked as outdated.

Sign in to view

alexeykudinkin force-pushed the ak/bndl-blk-fix branch from bad2cd7 to 65e7295 Compare October 27, 2025 18:07

alexeykudinkin mentioned this pull request Oct 27, 2025

[data] Assert internal + external queues are empty upon completion #57631

Merged

8 tasks

alexeykudinkin added 10 commits October 27, 2025 12:14

Added BundleQueue.num_blocks to track number of actual blocks in a …

9a8f793

…bundle Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Added InternalQueueOperatorMixin.internal_queue_size_bytes tracking…

82eb7c5

… num bytes held inside internal queue; Revisited `InternalQueueOperatorMixin.internal_queue_size` to remove assumption that every bundle holds just 1 block Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidiyng up

15c5264

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Cleaning up naming

962b471

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Cleaning up more naming

fd59a96

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up

335e751

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

a60281b

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixed total_enqueued_num_blocks to count blocks instead of bundles

64ca0fc

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing handling of _num_blocks

81b6a50

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Rebased Union onto FIFOBundleQueue

7c77c6a

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin added 10 commits October 27, 2025 12:14

Tidying up

36f7b1b

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

17530a0

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Reverting metric changes

a52e82c

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Tidying up;

4ed7dc9

Fixing tests Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing more

0de409c

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

6ece13f

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Refined naming

89e9c21

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Revert

cfc00d5

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

Fixing more tests

e8b1f37

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

lint

7165108

Signed-off-by: Alexey Kudinkin <ak@anyscale.com>

alexeykudinkin force-pushed the ak/bndl-blk-fix branch from 65e7295 to 7165108 Compare October 27, 2025 19:14

alexeykudinkin enabled auto-merge (squash) October 27, 2025 20:02

alexeykudinkin merged commit 7dd8eb2 into master Oct 27, 2025
7 checks passed

alexeykudinkin deleted the ak/bndl-blk-fix branch October 27, 2025 20:31

alexeykudinkin restored the ak/bndl-blk-fix branch October 27, 2025 20:37

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Data] Replace number of bundles with proper number of blocks#58030

[Data] Replace number of bundles with proper number of blocks#58030
alexeykudinkin merged 20 commits intomasterfrom
ak/bndl-blk-fix

alexeykudinkin commented Oct 23, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

bveeramani Oct 24, 2025

Uh oh!

iamjustinhsu Oct 24, 2025

Uh oh!

iamjustinhsu Oct 24, 2025

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		self._bundle_buffer_size = 0
		self._bundle_buffer_size_bytes = 0

Conversation

alexeykudinkin commented Oct 23, 2025

Description

Related issues

Additional information

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

bveeramani Oct 24, 2025

Choose a reason for hiding this comment

Uh oh!

iamjustinhsu Oct 24, 2025

Choose a reason for hiding this comment

Uh oh!

iamjustinhsu Oct 24, 2025

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants