[Data] fix cap_resource_request_limits to handle heterogenous nodes and refactor

Signed-off-by: Marwan Sarieddine <sarieddine.marwan@gmail.com>

Author: marwan116

python/ray/data/_internal/cluster_autoscaler/default_cluster_autoscaler.py

@@ -1,10 +1,10 @@
-import logging
 import math
 import time
-from typing import TYPE_CHECKING, Dict, List
+from typing import TYPE_CHECKING, Dict
 
 import ray
 from .base_cluster_autoscaler import ClusterAutoscaler
+from .util import cap_resource_request_to_limits
 from ray.data._internal.execution.autoscaling_requester import (
     get_or_create_autoscaling_requester_actor,
 )
@@ -14,9 +14,6 @@
     from ray.data._internal.execution.streaming_executor_state import Topology
 
 
-logger = logging.getLogger(__name__)
-
-
 class DefaultClusterAutoscaler(ClusterAutoscaler):
     # Min number of seconds between two autoscaling requests.
     MIN_GAP_BETWEEN_AUTOSCALING_REQUESTS = 20
@@ -35,46 +32,6 @@ def __init__(
         # Last time when a request was sent to Ray's autoscaler.
         self._last_request_time = 0
 
-    def _cap_resource_request_to_limits(
-        self, resource_request: List[Dict]
-    ) -> List[Dict]:
-        """Cap the resource request to not exceed user-configured resource limits.
-
-        If the user has set explicit (non-infinite) resource limits, this method
-        filters the resource request to ensure the total requested resources do not
-        exceed those limits.
-
-        Args:
-            resource_request: List of resource bundles to request.
-
-        Returns:
-            A filtered list of resource bundles that respects user limits.
-        """
-        limits = self._resource_limits
-
-        # If no explicit limits are set (all infinite), return the original request
-        if limits == ExecutionResources.inf():
-            return resource_request
-
-        capped_request = []
-        total = ExecutionResources.zero()
-
-        for bundle in resource_request:
-            new_total = total.add(ExecutionResources.from_resource_dict(bundle))
-
-            if not new_total.satisfies_limit(limits):
-                logger.debug(
-                    f"Capped autoscaling resource request from {len(resource_request)} "
-                    f"bundles to {len(capped_request)} bundles to respect "
-                    f"user-configured resource limits: {limits}."
-                )
-                break
-
-            capped_request.append(bundle)
-            total = new_total
-
-        return capped_request
-
     def try_trigger_scaling(self):
         """Try to scale up the cluster to accommodate the provided in-progress workload.
 
@@ -132,7 +89,9 @@ def to_bundle(resource: ExecutionResources) -> Dict:
                 resource_request.append(task_bundle)
 
         # Cap the resource request to respect user-configured limits
-        resource_request = self._cap_resource_request_to_limits(resource_request)
+        resource_request = cap_resource_request_to_limits(
+            resource_request, self._resource_limits
+        )
 
         self._send_resource_request(resource_request)
 

python/ray/data/_internal/cluster_autoscaler/default_cluster_autoscaler_v2.py

@@ -4,7 +4,7 @@
 from collections import defaultdict
 from dataclasses import dataclass
 from logging import getLogger
-from typing import TYPE_CHECKING, Callable, Dict, List, Optional
+from typing import TYPE_CHECKING, Callable, Dict, Optional
 
 import ray
 from .base_autoscaling_coordinator import AutoscalingCoordinator
@@ -15,6 +15,7 @@
     ResourceUtilizationGauge,
     RollingLogicalUtilizationGauge,
 )
+from .util import cap_resource_request_to_limits
 from ray.data._internal.cluster_autoscaler import ClusterAutoscaler
 from ray.data._internal.execution.interfaces.execution_options import ExecutionResources
 
@@ -163,46 +164,6 @@ def __init__(
         # so the first `get_total_resources` call can get the allocated resources.
         self._send_resource_request([])
 
-    def _cap_resource_request_to_limits(
-        self, resource_request: List[Dict]
-    ) -> List[Dict]:
-        """Cap the resource request to not exceed user-configured resource limits.
-
-        If the user has set explicit (non-infinite) resource limits, this method
-        filters the resource request to ensure the total requested resources do not
-        exceed those limits.
-
-        Args:
-            resource_request: List of resource bundles (node specs) to request.
-
-        Returns:
-            A filtered list of resource bundles that respects user limits.
-        """
-        limits = self._resource_limits
-
-        # If no explicit limits are set (all infinite), return the original request
-        if limits == ExecutionResources.inf():
-            return resource_request
-
-        capped_request = []
-        total = ExecutionResources.zero()
-
-        for bundle in resource_request:
-            new_total = total.add(ExecutionResources.from_resource_dict(bundle))
-
-            if not new_total.satisfies_limit(limits):
-                logger.debug(
-                    f"Capped autoscaling resource request from {len(resource_request)} "
-                    f"bundles to {len(capped_request)} bundles to respect "
-                    f"user-configured resource limits: {limits}."
-                )
-                break
-
-            capped_request.append(bundle)
-            total = new_total
-
-        return capped_request
-
     def try_trigger_scaling(self):
         # Note, should call this method before checking `_last_request_time`,
         # in order to update the average cluster utilization.
@@ -256,7 +217,9 @@ def try_trigger_scaling(self):
         logger.debug(debug_msg)
 
         # Cap the resource request to respect user-configured limits
-        resource_request = self._cap_resource_request_to_limits(resource_request)
+        resource_request = cap_resource_request_to_limits(
+            resource_request, self._resource_limits
+        )
 
         self._send_resource_request(resource_request)
 

python/ray/data/_internal/cluster_autoscaler/util.py

@@ -0,0 +1,66 @@
+import logging
+from typing import Dict, List
+
+from ray.data._internal.execution.interfaces import ExecutionResources
+
+logger = logging.getLogger(__name__)
+
+
+def cap_resource_request_to_limits(
+    resource_request: List[Dict],
+    resource_limits: ExecutionResources,
+) -> List[Dict]:
+    """Cap the resource request to not exceed user-configured resource limits.
+
+    If the user has set explicit (non-infinite) resource limits, this function
+    filters the resource request to ensure the total requested resources do not
+    exceed those limits.
+
+    Bundles are sorted by size (smallest first) to maximize the number of bundles
+    that can fit within the limits. This ensures that smaller bundles are not
+    excluded just because a larger bundle appeared earlier in iteration order.
+
+    Args:
+        resource_request: List of resource bundles to request.
+        resource_limits: The user-configured resource limits.
+
+    Returns:
+        A filtered list of resource bundles that respects user limits.
+    """
+    # If no explicit limits are set (all infinite), return the original request
+    if resource_limits == ExecutionResources.inf():
+        return resource_request
+
+    # Sort bundles by size (smallest first) to maximize packing within limits.
+    # This ensures smaller bundles aren't excluded due to larger bundles
+    # appearing earlier in arbitrary iteration order.
+    def bundle_sort_key(bundle: Dict) -> tuple:
+        return (
+            bundle.get("CPU", 0),
+            bundle.get("GPU", 0),
+            bundle.get("memory", 0),
+        )
+
+    sorted_bundles = sorted(resource_request, key=bundle_sort_key)
+
+    capped_request = []
+    total = ExecutionResources.zero()
+
+    for bundle in sorted_bundles:
+        new_total = total.add(ExecutionResources.from_resource_dict(bundle))
+
+        # Skip bundles that don't fit, continue checking smaller ones
+        if not new_total.satisfies_limit(resource_limits):
+            continue
+
+        capped_request.append(bundle)
+        total = new_total
+
+    if len(capped_request) < len(resource_request):
+        logger.debug(
+            f"Capped autoscaling resource request from {len(resource_request)} "
+            f"bundles to {len(capped_request)} bundles to respect "
+            f"user-configured resource limits: {resource_limits}."
+        )
+
+    return capped_request

python/ray/data/tests/test_default_cluster_autoscaler_v2.py

@@ -402,6 +402,57 @@ def test_try_scale_up_respects_resource_limits(
         assert resources_allocated.gpu == node_spec.gpu * expected_nodes
         assert resources_allocated.memory == node_spec.mem * expected_nodes
 
+    def test_try_scale_up_respects_resource_limits_heterogeneous_nodes(self):
+        """Test that smaller bundles are included even when larger bundles exceed limits.
+
+        This tests the fix for an issue where heterogeneous node types could result
+        in empty or suboptimal resource requests if a large bundle appeared first
+        in iteration order and exceeded limits, causing smaller valid bundles to
+        be skipped.
+        """
+        # Set a CPU limit that:
+        # - Is smaller than a single large node (12 CPUs)
+        # - But can fit multiple small nodes (4 CPUs each)
+        resource_limits = ExecutionResources(cpu=10)
+
+        large_node_spec = _NodeResourceSpec.of(cpu=12, gpu=1, mem=8000)
+        small_node_spec = _NodeResourceSpec.of(cpu=4, gpu=0, mem=2000)
+
+        scale_up_threshold = 0.75
+        utilization = ExecutionResources(cpu=0.9, gpu=0.9, object_store_memory=0.9)
+        fake_coordinator = FakeAutoscalingCoordinator()
+
+        # Return heterogeneous node types - the order here shouldn't matter
+        # because the implementation should sort bundles by size
+        def get_heterogeneous_nodes():
+            return {
+                large_node_spec: 1,  # 1 existing large node, wants 2 bundles
+                small_node_spec: 1,  # 1 existing small node, wants 2 bundles
+            }
+
+        autoscaler = DefaultClusterAutoscalerV2(
+            resource_manager=MagicMock(),
+            resource_limits=resource_limits,
+            execution_id="test_execution_id",
+            cluster_scaling_up_delta=1,
+            resource_utilization_calculator=StubUtilizationGauge(utilization),
+            cluster_scaling_up_util_threshold=scale_up_threshold,
+            min_gap_between_autoscaling_requests_s=0,
+            autoscaling_coordinator=fake_coordinator,
+            get_node_counts=get_heterogeneous_nodes,
+        )
+
+        autoscaler.try_trigger_scaling()
+
+        resources_allocated = autoscaler.get_total_resources()
+        # Should get 2 small nodes (8 CPUs) since large nodes (12 CPUs) exceed limit
+        assert resources_allocated.cpu == 8, (
+            f"Expected 8 CPUs (2 small nodes), got {resources_allocated.cpu}. "
+            "Smaller bundles should be included even when larger ones exceed limits."
+        )
+        assert resources_allocated.gpu == 0
+        assert resources_allocated.memory == 4000
+
 
 if __name__ == "__main__":
     import sys