wehub-resource-sync
98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
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
|