[Data] Incrementally update resource usage after task dispatch

[raygao25]

Description

This PR updates the Ray Data streaming executor to avoid doing a full resource usage recomputation after every task dispatch.
Previously, each dispatch in the scheduling loop called ResourceManager.update_usages(), which walks the full topology and recomputes usage for every operator. This can be expensive in the hot path. With this change, after dispatching a task, we only refresh usage for the dispatched operator and its immediate upstream operators, since those are the operators whose accounting can change from moving an input bundle into the operator.
Allocator budgets are intentionally left frozen during the inner dispatch loop to keep the hot path cheap. They are refreshed by the next full update_usages() call before the next scheduling-loop step.

Key changes:

• Adding ResourceManager.update_usages_incremental() to update per-op and global resource usage by delta.
• Updating the streaming executor dispatch loop to call update_usages_incremental instead of update_usages.
• Adding a unit test that compares incremental refresh against a full recompute.

Related issues

Additional information

TODO: test was done with 2.55.0-release. Need to test with latest main

[gemini-code-assist]

Code Review

This pull request introduces an incremental resource usage update mechanism in the ResourceManager to improve efficiency by targeting only affected operators and their upstreams. This change is integrated into the StreamingExecutor and verified with a new test case. Review feedback identifies a bug where an immutable resource addition was discarded, notes variable shadowing of the 'op' parameter, and suggests refactoring duplicated logic into a shared helper method to improve maintainability.

[gemini-code-assist]

The ExecutionResources class is immutable, and its add method returns a new instance rather than modifying the object in place. Consequently, the result of op_usage.add(...) is currently discarded, and the extra resource usage is not reflected in the operator's accounting or the global aggregates. This appears to be a pre-existing bug in update_usages that was duplicated here.

Suggested change

	op_usage.add(op.extra_resource_usage())
	op_usage = op_usage.add(op.extra_resource_usage())

[gemini-code-assist]

The loop variable op shadows the function parameter op. While this doesn't break the current logic since the parameter is not used after the loop starts, it is a poor practice that can lead to confusion or bugs during future maintenance. It's better to use a distinct name like affected_op for the iteration.

        for affected_op in affected_ops:
            state = self._topology[affected_op]

            # Subtract this op's previous contribution from globals.
            old_usage = self._op_usages.get(affected_op)
            old_running = self._op_running_usages.get(affected_op)
            old_pending = self._op_pending_usages.get(affected_op)
            if old_usage is not None:
                self._global_usage = self._global_usage.subtract(old_usage)
            if old_running is not None:
                self._global_running_usage = self._global_running_usage.subtract(
                    old_running
                )
            if old_pending is not None:
                self._global_pending_usage = self._global_pending_usage.subtract(
                    old_pending
                )

            op_usage = affected_op.current_logical_usage()
            op_running_usage = affected_op.running_logical_usage()
            op_pending_usage = affected_op.pending_logical_usage()

            assert not op_usage.object_store_memory
            assert not op_running_usage.object_store_memory
            assert not op_pending_usage.object_store_memory

            used_object_store = self._estimate_object_store_memory_usage(affected_op, state)

            op_usage = op_usage.copy(object_store_memory=used_object_store)
            op_running_usage = op_running_usage.copy(
                object_store_memory=used_object_store
            )

            if isinstance(affected_op, ReportsExtraResourceUsage):
                op_usage = op_usage.add(affected_op.extra_resource_usage())

            self._op_usages[affected_op] = op_usage
            self._op_running_usages[affected_op] = op_running_usage
            self._op_pending_usages[affected_op] = op_pending_usage

            self._global_usage = self._global_usage.add(op_usage)
            self._global_running_usage = self._global_running_usage.add(
                op_running_usage
            )
            self._global_pending_usage = self._global_pending_usage.add(
                op_pending_usage
            )

            affected_op._metrics.obj_store_mem_used = op_usage.object_store_memory

[gemini-code-assist]

The logic for recomputing an operator's resource usage and updating its associated metrics and caches is duplicated between update_usages (lines 227-261) and update_usages_incremental (lines 298-328). Consider extracting this logic into a private helper method, e.g., _recompute_op_usage(self, op: PhysicalOperator), to improve maintainability and ensure consistency.

[cursor]

Loop variable shadows the function parameter op

Low Severity

The for op in affected_ops: loop variable shadows the function parameter op of update_usages_incremental. While not currently causing a functional bug (since the parameter is fully consumed before the loop begins), this makes the code fragile — any future modification that references the parameter op after the loop would silently get the last element from the set iteration instead, leading to a hard-to-diagnose bug.

 

^{Reviewed by Cursor Bugbot for commit 484e689e27e091eaa7dc411389bef6d67d85f7a0. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{Reviewed by Cursor Bugbot for commit dc24086. Configure here.}

[cursor]

Discarded return value loses extra resource usage

Medium Severity

op_usage.add(op.extra_resource_usage()) discards its return value. ExecutionResources.add() returns a new object rather than mutating in place, so the extra resource usage from ReportsExtraResourceUsage operators is never included in op_usage. The value stored in self._op_usages[op] and added to global aggregates will be missing this contribution. The same pattern exists in update_usages() line 243, so both methods are equally broken and the test passes, but extra resource usage is silently dropped.

 

^{Reviewed by Cursor Bugbot for commit dc24086. Configure here.}

[iamjustinhsu]

Thanks for taking this on! I left a comment about the redesign this in general to avoid the complexity of maintaining both update_usages and update_usages_incremental

[iamjustinhsu]

Hmm, while this probably does work, I think having 2 separate update usages might be confusing long-term. If ur up for it, this is what I recommend:

1. Split the current update_usages into 2 functions, one called update_budgets, one called update_usages_incremental.
2. In update_budgets, we will call self._update_allocated_budgets()
3. In update_usages_incremental, we'll perform a delta of the current usages. One for task dispatch (+), another for task return (-1)
4. Then in _scheduling_loop_step, we will do these steps
   • process_completed_tasks
     • update_usages_incremental(op) for each operator task that returns
   • topology[op].dispatch_next_tasks + update_usages_incremental(op)

[iamjustinhsu]

Can you help me understand why this is needed? I would assume that that if we are updating the usages for operator, it wouldn't affect the usages of the immediate upstream operator?

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[github-actions]

This pull request has been automatically closed because there has been no more activity in the 14 days
since being marked stale.

Please feel free to reopen or open a new pull request if you'd still like this to be addressed.

Again, you can always ask for help on our discussion forum or Ray's public slack channel.

Thanks again for your contribution!