[Data][Flaky] Ensure ActorPoolMapOperator clears all queues on completion

[400Ping]

Description

One of the PyTorch examples intermittently fails with an error like this:

AssertionError: Expected Internal Input Queue for MapBatches(ResnetModel) to be empty, but found 1 bundles

This assertion can fail because of two bugs:

1. We don't call done_adding_bundles before clearing the block ref bundler
2. We don't clean the actor pool's bundle queue

This PR fixes those bugs and deflakes the test.

Related issues

Closes #58546

Additional information

[gemini-code-assist]

Code Review

This pull request addresses a flaky test by fixing a race condition in ActorPoolMapOperator. The change adds a call to _dispatch_tasks() in all_inputs_done() to ensure any queued bundles are processed when no more inputs are expected. This is a solid fix for the described problem. I've added one minor suggestion to clean up some redundant code in the same method for improved maintainability.

[400Ping]

cc @owenowenisme

[iamjustinhsu]

Hi @400Ping, thanks for your contribution! We really appreciate the effort.

I think @owenowenisme might be right here, all_inputs_done should already be called. My hunch is that BlockRefBundler class is not correctly finalizing the outputs. From here

            if (
                output_buffer_size < self._min_rows_per_bundle
                or output_buffer_size == 0
            ):

it looks like we should be doing

            if (
                output_buffer_size < self._min_rows_per_bundle
                or output_buffer_size == 0
                or self._finalized
            ):

so that there are no more remainders. Can you try that and report back?
cc: @bveeramani

[bveeramani]

Hey @400Ping, would you mind helping me understand the root cause of the assertion error?

Also, did the repro script fail before the changes?

[400Ping]

Hey @400Ping, would you mind helping me understand the root cause of the assertion error?

Also, did the repro script fail before the changes?

Ok, will try to find it.

[iamjustinhsu]

@bveeramani this feels like this break should have been there in the beginning, right?

[400Ping]

cc @bveeramani PTAL

[400Ping]

not very sure if this is the way to solve this.

[iamjustinhsu]

@bveeramani my impression is that when self._finalized=True (ie, when an operator is completed()), it is possible this for loop enters the else statement down below, populating remainders with non-empty ref bundles.

I also think the break statement is necessary too, otherwise, remainders is always being reassigned.

[bveeramani]

@400Ping @iamjustinhsu I was able to create a minimal repro of this issue (at least, I think this is the same issue)

import ray


class Fn:

    def __call__(self, batch):
        return batch


ds = ray.data.range(100, override_num_blocks=100).map_batches(Fn, batch_size=10).limit(1)
for _ in ds.iter_internal_ref_bundles():
    pass

Traceback (most recent call last):                                                                                                                                      
  File "/Users/balaji/ray/1.py", line 12, in <module>
    for _ in ds.iter_internal_ref_bundles():                                                                                                                            
  File "/Users/balaji/ray/python/ray/data/_internal/execution/interfaces/executor.py", line 34, in __next__
    return self.get_next()
  File "/Users/balaji/ray/python/ray/data/_internal/execution/legacy_compat.py", line 76, in get_next
    bundle = self._base_iterator.get_next(output_split_idx)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 786, in get_next
    bundle = state.get_output_blocking(output_split_idx)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor_state.py", line 454, in get_output_blocking
    raise self._exception
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 356, in run
    continue_sched = self._scheduling_loop_step(self._topology)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 529, in _scheduling_loop_step
    self._validate_operator_queues_empty(op, state)
  File "/Users/balaji/ray/python/ray/data/_internal/execution/streaming_executor.py", line 571, in _validate_operator_queues_empty
    assert op.internal_input_queue_num_blocks() == 0, error_msg.format(
AssertionError: Expected Internal Input Queue for MapBatches(Fn) to be empty, but found 8 bundles

@iamjustinhsu would a reasonable fix be to just clear the internal queues when the operator is manually marked finished?

    def mark_execution_finished(self):
        # Discard remaining bundles in the internal bundle queue.
        self._bundle_queue.clear()

        # Discard remaining bundles in the block ref bundler.
        self._block_ref_bundler.done_adding_bundles()
        while self._block_ref_bundler.has_bundle():
            self._block_ref_bundler.get_next_bundle()

        super().mark_execution_finished()

[iamjustinhsu]

@bveeramani oh i see now. We actually do this for for all InternalQueueOperatorMixin operators. However, I implemented this for MapOperator, thinking ActorPoolMapOperator uses the same underlying implementation. It actually doesn't, and contains an additional queue called _bundle_queue. Here is my PR for reference: #58441. What we should do is have logic for also clearing the self._bundle_queue in ActorPoolMapOperator. This can be done with having separate implementations for TaskPoolMapOperator and ActorPoolMapOperator

[bveeramani]

Ah, okay.

@400Ping would you mind refactoring the PR or opening a new PR to do the following:

• Delete MapOperator.clear_internal_output_queue and MapOperator.clear_internal_input_queue
• Add TaskPoolMapOperator.clear_internal_output_queue and TaskPoolMapOperator.clear_internal_input_queue
• Add ActorPoolMapOperator.clear_internal_output_queue and ActorPoolMapOperator.clear_internal_input_queue

I think this is all that we need to do to solve the flaky issue.

[bveeramani]

The implementations for clear_internal_input_queue and clear_internal_output_queue LGTM.

@400Ping to keep the git history clear, would you mind reverting all of the changes unrelated to those methods (e.g., formatting, removing _inputs_done, type annotations)? Once we've reverted the unrelated changes, I'll approve the PR

[400Ping]

The implementations for clear_internal_input_queue and clear_internal_output_queue LGTM.

@400Ping to keep the git history clear, would you mind reverting all of the changes unrelated to those methods (e.g., formatting, removing _inputs_done, type annotations)? Once we've reverted the unrelated changes, I'll approve the PR

Ok.

[bveeramani]

LGTM! 🚢

[400Ping]

Thanks for the fix!

[bveeramani]

@400Ping Just merged! ty for the contribution!

[400Ping]

@400Ping Just merged! ty for the contribution!

Thank you as well, I am a newbie in this area 😓.