[Data] Fix streaming executor to drain upstream output queue(s)

[srinathk10]

Why are these changes needed?

Issue

• Consider a Pipeline ReadFiles->Limit->Map->...
• When Limit, we mark the mark_execution_finished for current
  LimitOperator, so it won't dequeue from the upstream Ops output queue
  anymore.
• When this happens, we can have Object Store resource leaks for the
  life of the pipeline because of pending unneeded Blocks in the upstream
  Op's output queue.
• In the below progress bar, ReadFiles has 9.7GB locked up.

(pid=18548) ✔️  Dataset train_14_0 execution finished in 130.77 seconds: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 887k/887k [02:10<00:00, 6.78k row/s] 
(pid=18548) - ListFiles: Tasks: 0; Actors: 0; Queued blocks: 0; Resources: 0.0 CPU, 0.0B object store (in=0.0B,out=0.0B): 100%|███████████████████████████████████████████████████████████████████████████████████████████████████| 513/513 [02:10<00:00, 3.92 row/s]
(pid=18548) - ReadFiles: Tasks: 0; Actors: 0; Queued blocks: 0; Resources: 0.0 CPU, 9.7GB object store (in=0.0B,out=9.7GB): 100%|████████████████████████████████████████████████████████████████████████████████████████████| 1.28M/1.28M [02:10<00:00, 9.79k row/s]
(pid=18548) - limit=1000000: Tasks: 0; Actors: 0; Queued blocks: 107; Resources: 0.0 CPU, 0.0B object store (in=0.0B,out=0.0B): 100%|████████████████████████████████████████████████████████████████████████████████████████| 1.00M/1.00M [02:10<00:00, 7.64k row/s]
(pid=18548) - Map(map_fn): Tasks: 0; Actors: 0; Queued blocks: 0; Resources: 0.0 CPU, 0.0B object store (in=0.0B,out=0.0B): 100%|████████████████████████████████████████████████████████████████████████████████████████████| 1.00M/1.00M [02:10<00:00, 7.64k row/s]
(pid=18548) - split(16, equal=True): Tasks: 0; Actors: 0; Queued blocks: 0; Resources: 0.0 CPU, 63.7GB object store (in=0.0B,out=63.7GB); [locality disabled]: 100%|█████████████████████████████████████████████████████████| 1.00M/1.00M [02:10<00:00, 7.64k row/s]

Fix

• When current Op is done, proactively drain the upstream Ops output
  queues and mark input done for that upstream Op.
• Here you can observe that ReadFiles Op does not have object store
  resources locked up anymore.

Related issue number

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

[gemini-code-assist]

Code Review

This pull request effectively addresses a potential memory leak in the streaming executor by ensuring upstream output queues are drained when a downstream operator finishes execution, such as when a limit is reached. The fix in update_operator_states is correct and is accompanied by two new tests: a focused unit test and a valuable integration test that validates the fix in an end-to-end scenario. My review found a minor issue in the new integration test's logic for identifying an operator, which could make the test fragile. Overall, this is a solid contribution that improves resource management in Ray Data pipelines.

[gemini-code-assist]

The condition "ReadParquet" in op.name is likely incorrect and fragile. The name of the logical Read operator for a Parquet datasource is constructed as f"Read({datasource.get_name()})", which results in Read(Parquet). The substring "ReadParquet" is not present in this name. A more robust approach would be to match the exact name.

Suggested change

	if "ReadParquet" in op.name:
	if "Read(Parquet)" in op.name:

[cursor]

Bug: Downstream Completion Triggers Premature Queue Clearing

Unconditionally clearing an upstream operator's output queue when any downstream operator finishes execution can lead to data loss. In fan-out scenarios, this starves other active downstream operators that still depend on that data.