Open
Conversation
Signed-off-by: will <zzchun8@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a resource-based actor autoscaler and a pluggable actor removal strategy, which are significant enhancements for managing resources in Ray Data. The implementation is well-structured and includes comprehensive tests for the new functionality.
My review has identified a couple of areas for improvement:
- There is some duplicated code in the new
ResourceBasedActorAutoscalerthat could be refactored for better maintainability. - A constant appears to be defined twice in
DataContext, which should be corrected.
Overall, this is a solid contribution. Addressing these points will further improve the code quality.
Comment on lines
+226
to
+228
| DEFAULT_ACTOR_MAX_TASKS_IN_FLIGHT_TO_MAX_CONCURRENCY_FACTOR = env_integer( | ||
| "RAY_DATA_ACTOR_DEFAULT_MAX_TASKS_IN_FLIGHT_TO_MAX_CONCURRENCY_FACTOR", 2 | ||
| ) |
There was a problem hiding this comment.
This constant DEFAULT_ACTOR_MAX_TASKS_IN_FLIGHT_TO_MAX_CONCURRENCY_FACTOR appears to be a duplicate. It is already defined on line 217. Please remove this redundant definition.
Comment on lines
+157
to
+246
| def _calculate_pool_min_resources( | ||
| self, | ||
| actor_pool: "AutoscalingActorPool", | ||
| per_actor_resources: ExecutionResources, | ||
| weight_ratio: float, | ||
| all_pools: list, | ||
| pool_weights: Dict, | ||
| pool_resource_needs: Dict, | ||
| ) -> ExecutionResources: | ||
| """Calculate min resources for a single pool""" | ||
| pool_min_cpu = ( | ||
| self._job_min_resources.cpu * weight_ratio | ||
| if per_actor_resources.cpu > 0 | ||
| else 0 | ||
| ) | ||
|
|
||
| # GPU allocation: Only allocate to pools that need GPUs | ||
| if per_actor_resources.gpu > 0: | ||
| gpu_pools = [p for p in all_pools if pool_resource_needs[p].gpu > 0] | ||
| if gpu_pools: | ||
| gpu_total_weight = sum(pool_weights[p] for p in gpu_pools) | ||
| gpu_weight_ratio = ( | ||
| pool_weights[actor_pool] / gpu_total_weight | ||
| if gpu_total_weight > 0 | ||
| else 1.0 / len(gpu_pools) | ||
| ) | ||
| pool_min_gpu = self._job_min_resources.gpu * gpu_weight_ratio | ||
| else: | ||
| pool_min_gpu = 0 | ||
| else: | ||
| pool_min_gpu = 0 | ||
|
|
||
| # Memory allocation | ||
| pool_min_memory = ( | ||
| self._job_min_resources.memory * weight_ratio | ||
| if per_actor_resources.memory > 0 | ||
| else 0 | ||
| ) | ||
|
|
||
| return ExecutionResources( | ||
| cpu=pool_min_cpu, | ||
| gpu=pool_min_gpu, | ||
| memory=pool_min_memory, | ||
| ) | ||
|
|
||
| def _calculate_pool_max_resources( | ||
| self, | ||
| actor_pool: "AutoscalingActorPool", | ||
| per_actor_resources: ExecutionResources, | ||
| weight_ratio: float, | ||
| all_pools: list, | ||
| pool_weights: Dict, | ||
| pool_resource_needs: Dict, | ||
| ) -> ExecutionResources: | ||
| """Calculate max resources for a single pool""" | ||
| # CPU allocation | ||
| pool_max_cpu = ( | ||
| self._job_max_resources.cpu * weight_ratio | ||
| if per_actor_resources.cpu > 0 | ||
| else 0 | ||
| ) | ||
|
|
||
| # GPU allocation: Only allocate to pools that need GPUs | ||
| if per_actor_resources.gpu > 0: | ||
| gpu_pools = [p for p in all_pools if pool_resource_needs[p].gpu > 0] | ||
| if gpu_pools: | ||
| gpu_total_weight = sum(pool_weights[p] for p in gpu_pools) | ||
| gpu_weight_ratio = ( | ||
| pool_weights[actor_pool] / gpu_total_weight | ||
| if gpu_total_weight > 0 | ||
| else 1.0 / len(gpu_pools) | ||
| ) | ||
| pool_max_gpu = self._job_max_resources.gpu * gpu_weight_ratio | ||
| else: | ||
| pool_max_gpu = 0 | ||
| else: | ||
| pool_max_gpu = 0 | ||
|
|
||
| # Memory allocation | ||
| pool_max_memory = ( | ||
| self._job_max_resources.memory * weight_ratio | ||
| if per_actor_resources.memory > 0 | ||
| else 0 | ||
| ) | ||
|
|
||
| return ExecutionResources( | ||
| cpu=pool_max_cpu, | ||
| gpu=pool_max_gpu, | ||
| memory=pool_max_memory, | ||
| ) |
There was a problem hiding this comment.
The methods _calculate_pool_min_resources and _calculate_pool_max_resources contain a lot of duplicated logic. To improve maintainability and reduce redundancy, you can extract the common logic into a single helper method that takes the job-level resources (min or max) as an argument.
Here's a suggested refactoring:
def _calculate_pool_resources(
self,
job_resources: ExecutionResources,
actor_pool: "AutoscalingActorPool",
per_actor_resources: ExecutionResources,
weight_ratio: float,
all_pools: list,
pool_weights: Dict,
pool_resource_needs: Dict,
) -> ExecutionResources:
"""Calculate resources for a single pool based on job-level resources."""
pool_cpu = (
job_resources.cpu * weight_ratio
if per_actor_resources.cpu > 0
else 0
)
# GPU allocation: Only allocate to pools that need GPUs
if per_actor_resources.gpu > 0:
gpu_pools = [p for p in all_pools if pool_resource_needs[p].gpu > 0]
if gpu_pools:
gpu_total_weight = sum(pool_weights[p] for p in gpu_pools)
gpu_weight_ratio = (
pool_weights[actor_pool] / gpu_total_weight
if gpu_total_weight > 0
else 1.0 / len(gpu_pools)
)
pool_gpu = job_resources.gpu * gpu_weight_ratio
else:
pool_gpu = 0
else:
pool_gpu = 0
# Memory allocation
pool_memory = (
job_resources.memory * weight_ratio
if per_actor_resources.memory > 0
else 0
)
return ExecutionResources(
cpu=pool_cpu,
gpu=pool_gpu,
memory=pool_memory,
)
def _calculate_pool_min_resources(
self,
actor_pool: "AutoscalingActorPool",
per_actor_resources: ExecutionResources,
weight_ratio: float,
all_pools: list,
pool_weights: Dict,
pool_resource_needs: Dict,
) -> ExecutionResources:
"""Calculate min resources for a single pool"""
return self._calculate_pool_resources(
self._job_min_resources,
actor_pool,
per_actor_resources,
weight_ratio,
all_pools,
pool_weights,
pool_resource_needs,
)
def _calculate_pool_max_resources(
self,
actor_pool: "AutoscalingActorPool",
per_actor_resources: ExecutionResources,
weight_ratio: float,
all_pools: list,
pool_weights: Dict,
pool_resource_needs: Dict,
) -> ExecutionResources:
"""Calculate max resources for a single pool"""
return self._calculate_pool_resources(
self._job_max_resources,
actor_pool,
per_actor_resources,
weight_ratio,
all_pools,
pool_weights,
pool_resource_needs,
)
ffbin
approved these changes
Jan 16, 2026
ffbin
reviewed
Jan 16, 2026
| logger = logging.getLogger(__name__) | ||
|
|
||
| # Default maximum pool size | ||
| DEFAULT_MAX_POOL_SIZE = 100 |
|
|
||
| logger = logging.getLogger(__name__) | ||
|
|
||
| # Default maximum pool size |
Collaborator
There was a problem hiding this comment.
# Default maximum pool size.
| config: AutoscalingConfig, | ||
| ): | ||
| super().__init__(topology, resource_manager, config=config) | ||
| # job-level resource limits |
Collaborator
There was a problem hiding this comment.
# Job-level resource limits.
| """ | ||
|
|
||
| if min_resources is not None and max_resources is not None: | ||
| # Check CPU |
|
This pull request has been automatically marked as stale because it has not had 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. |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
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.
Description
This PR implements a resource-based actor autoscaler to improve resource management in Ray Data.
Features
Main Changes