[data] Allocate GPU resources in ResourceManager #54445
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Hao Chen <[email protected]>
Signed-off-by: Hao Chen <[email protected]>
Signed-off-by: Hao Chen <[email protected]>
Signed-off-by: Hao Chen <[email protected]>
Signed-off-by: Hao Chen <[email protected]>
Signed-off-by: Hao Chen <[email protected]>
There was a problem hiding this comment.
Pull Request Overview
Adds GPU-specific budgeting to the ResourceManager so that operators requesting GPUs receive a slice of the cluster’s GPUs rather than an infinite budget.
- Change GPU budget in
update_usagesto allocate “global GPUs minus current usage” for GPU operators; non-GPU operators get zero. - Update existing tests to expect a zero GPU budget instead of
infin memory-only cases. - Add
test_gpu_allocationandtest_multiple_gpu_operatorsto verify GPU allocations.
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| python/ray/data/_internal/execution/resource_manager.py | Implement per-op GPU budgeting in update_usages |
| python/ray/data/tests/test_resource_manager.py | Adjust memory-budget assertions and add GPU allocation tests |
Comments suppressed due to low confidence (2)
python/ray/data/tests/test_resource_manager.py:672
- [nitpick] Add a test case where an operator’s current GPU usage exceeds the global GPU limit to verify that the allocator properly clamps the budget to zero.
def test_multiple_gpu_operators(self, restore_data_context):
python/ray/data/_internal/execution/resource_manager.py:701
- Cache the result of
op.min_max_resource_requirements()in a local variable to avoid calling the method twice and improve readability.
if op.min_max_resource_requirements()[1].gpu > 0:
Signed-off-by: Hao Chen <[email protected]>
Signed-off-by: Hao Chen <[email protected]>
Signed-off-by: Hao Chen <[email protected]>
|
this PR also includes a few small drive-by fixes left from #54376 cc @alexeykudinkin |
bveeramani
reviewed
Jul 11, 2025
Comment on lines
+454
to
+461
| max_bytes_to_read = min( | ||
| ( | ||
| limit | ||
| for policy in backpressure_policies | ||
| if (limit := policy.max_task_output_bytes_to_read(op)) is not None | ||
| ), | ||
| default=None, | ||
| ) |
bveeramani
approved these changes
Jul 11, 2025
Comment on lines
+629
to
+630
| DataContext.get_current().op_resource_reservation_enabled = True | ||
| DataContext.get_current().op_resource_reservation_ratio = 0.5 |
There was a problem hiding this comment.
Here and for the test below -- aren't these the defaults? Are they necessary for this test?
There was a problem hiding this comment.
I'd like to make the test not depend on the defaults.
so it doesn't break if we change the behavior. (We just need to test the logic. defaults don't matter)
Comment on lines
+659
to
+660
| resource_manager._mem_op_internal = dict.fromkeys([o1, o2, o3], 0) | ||
| resource_manager._mem_op_outputs = dict.fromkeys([o1, o2, o3], 0) |
There was a problem hiding this comment.
OOC why do we need to configure these?
There was a problem hiding this comment.
actually not needed. removed
Signed-off-by: Hao Chen <[email protected]>
alimaazamat
pushed a commit
to alimaazamat/ray
that referenced
this pull request
Jul 24, 2025
Allocate GPU resources in ResourceManager. Currently we just allocate all available GPUs to all operators that need GPUs. If you have multiple GPU ops, each of them will get all GPUs. This PR is mainly to make the resource budget reporting correct. --------- Signed-off-by: Hao Chen <[email protected]> Signed-off-by: alimaazamat <[email protected]>
jugalshah291
pushed a commit
to jugalshah291/ray_fork
that referenced
this pull request
Sep 11, 2025
Allocate GPU resources in ResourceManager. Currently we just allocate all available GPUs to all operators that need GPUs. If you have multiple GPU ops, each of them will get all GPUs. This PR is mainly to make the resource budget reporting correct. --------- Signed-off-by: Hao Chen <[email protected]> Signed-off-by: jugalshah291 <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Allocate GPU resources in ResourceManager.
Currently we just allocate all available GPUs to all operators that need GPUs. If you have multiple GPU ops, each of them will get all GPUs.
This PR is mainly to make the resource budget reporting correct.