[Data] - Don't reserve GPU budget for non-GPU tasks

[goutamvenkat-anyscale]

Description

Follow up for this PR: #59632 (comment)

Only assign GPU budget if the operator requires it.

Image classification Release Test: https://buildkite.com/ray-project/release/builds/73917#019b90b1-2a29-424f-861b-8715909fe02e

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.

[goutamvenkat-anyscale]

/gemini summary

[gemini-code-assist]

Code Review

This pull request introduces a sensible change to prevent non-GPU operators from reserving GPU budget, which is particularly important for unbounded operators. The changes in ActorPoolMapOperator and TaskPoolMapOperator correctly adjust the maximum resource requirements based on whether the operator uses GPUs. A comprehensive test case has been added to validate this behavior in a realistic inference pipeline, ensuring that GPU resources are allocated efficiently. The implementation is clean and the test is thorough. I've added one minor suggestion to refactor the test setup for better readability.

[gemini-code-assist]

Summary of Changes

This pull request refines Ray Data's resource allocation mechanism to prevent non-GPU tasks from being assigned GPU budget. By explicitly capping GPU requirements to zero for CPU-only operators, the system ensures that valuable GPU resources are exclusively available for tasks that genuinely need them, leading to improved resource efficiency and better performance for mixed-workload execution.

Highlights

• GPU Budget Allocation Logic: Modified the resource allocation logic for ActorPoolMapOperator and TaskPoolMapOperator to explicitly set GPU budget to zero for operators that do not require GPUs, even when using infinite resource limits.
• Resource Hoarding Prevention: This change prevents non-GPU tasks from inadvertently 'hoarding' or reserving GPU resources, ensuring more efficient utilization of GPU hardware.
• New Test Case: Introduced a comprehensive test case (test_gpu_not_reserved_for_non_gpu_operators) to validate that GPU resources are correctly reserved only for GPU-dependent operators in a simulated inference pipeline.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• python/ray/data/_internal/execution/operators/actor_pool_map_operator.py
  • Updated min_max_resource_requirements to set gpu=0 in ExecutionResources.for_limits() if num_gpus_per_actor is zero, preventing non-GPU actors from reserving GPU budget.
• python/ray/data/_internal/execution/operators/task_pool_map_operator.py
  • Updated min_max_resource_requirements to set gpu=0 in ExecutionResources.for_limits() if per_task.gpu is zero, preventing non-GPU tasks from reserving GPU budget.
• python/ray/data/tests/test_reservation_based_resource_allocator.py
  • Added test_gpu_not_reserved_for_non_gpu_operators to verify that CPU-only operators (e.g., Read, Write) are allocated zero GPUs, while GPU-dependent operators (e.g., Infer1, Infer2) correctly reserve GPUs in a pipeline.

Activity

• goutamvenkat-anyscale requested a summary of the pull request.
• gemini-code-assist[bot] provided a suggestion to refactor the new test case by defining the list of operators once to improve readability and avoid repetition.

[bveeramani]

LGTM except comment about CPUs

[bveeramani]

I think this implementation special-cases GPUs because it assumes that all tasks/actors require logical CPUs and memory, but I don't think that assumption holds.

For example, here's a common thing users do:

ds.map_batches(Inference, num_gpus=1, batch_size=...)

In this case, I don't think the Inference actors request any logical CPUs.

For this reason, should we also include CPUs and memory?

[alexeykudinkin]

Also simplify this conditional to be just:

  gpu=0 if num_gpus == 0 else max_actors * num_gpus

[bveeramani]

Sort of out-of-scope for this PR, but this might cause issues if users (or optimization rules like ConfigureMapTaskMemoryRule) specify ray_remote_args_fn.

For example, if a user does:

ds.map_batches(..., ray_remote_args_fn=lambda: {"num_cpus": 10})

Then the max resource usage will be num_cpus=1, but each task requires 10 CPUs.

An easy fix might be to consider ray_remote_args_fn when returning incremental_resource_usage