chore: import upstream snapshot with attribution

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

@@ -0,0 +1,97 @@
+import logging
+from typing import Dict, List
+
+from ray.data._internal.execution.interfaces import ExecutionResources
+
+logger = logging.getLogger(__name__)
+
+
+def is_autoscaling_enabled() -> bool:
+    """Check if any node type has autoscaling enabled (can scale up).
+
+    A node type is autoscalable if max_count == -1 (unlimited) or
+    max_count > min_count. If no cluster config is available or no node type
+    is autoscalable, returns False.
+    """
+    import ray._private.state
+
+    cluster_config = ray._private.state.state.get_cluster_config()
+    if not cluster_config or not cluster_config.node_group_configs:
+        return False
+    return any(
+        ngc.max_count == -1 or ngc.max_count > ngc.min_count
+        for ngc in cluster_config.node_group_configs
+        if ngc.resources and ngc.max_count != 0
+    )
+
+
+def cap_resource_request_to_limits(
+    active_bundles: List[Dict],
+    pending_bundles: List[Dict],
+    resource_limits: ExecutionResources,
+) -> List[Dict]:
+    """Cap the resource request to not exceed user-configured resource limits.
+
+    Active bundles (for running tasks or existing nodes) are always included first
+    since they represent resources already in use. Pending bundles (for future work
+    or scale-up requests) are then added best-effort, sorted smallest-first to
+    maximize packing within limits.
+
+    This ensures that resources for already-running tasks are never crowded out
+    by pending work from smaller operators.
+
+    Args:
+        active_bundles: Bundles for already-running tasks or existing nodes
+            (must include - these represent current resource usage).
+        pending_bundles: Bundles for pending work or scale-up requests
+            (best-effort - only added if within limits).
+        resource_limits: The user-configured resource limits.
+
+    Returns:
+        A list of resource bundles that respects user limits, with active bundles
+        always included first.
+    """
+    # If no explicit limits are set (all infinite), return everything
+    if resource_limits == ExecutionResources.inf():
+        return active_bundles + pending_bundles
+
+    # Always include active bundles first - they're already running/allocated
+    capped_request = list(active_bundles)
+    total = ExecutionResources.zero()
+    for bundle in active_bundles:
+        total = total.add(ExecutionResources.from_resource_dict(bundle))
+
+    # Sort pending bundles by size (smallest first) to maximize packing.
+    # 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_pending = sorted(pending_bundles, key=bundle_sort_key)
+
+    for bundle in sorted_pending:
+        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
+
+    total_input = len(active_bundles) + len(pending_bundles)
+    if len(capped_request) < total_input:
+        logger.debug(
+            f"Capped autoscaling resource request from {total_input} "
+            f"bundles to {len(capped_request)} bundles to respect "
+            f"user-configured resource limits: {resource_limits}. "
+            f"({len(active_bundles)} active bundles kept, "
+            f"{len(capped_request) - len(active_bundles)}/{len(pending_bundles)} "
+            f"pending bundles included)."
+        )
+
+    return capped_request