import logging import math import time from abc import ABC, abstractmethod from collections import defaultdict from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Optional from ray._common.utils import env_bool, env_float from ray.data._internal.execution import create_resource_allocator from ray.data._internal.execution.block_ref_counter import BlockRefCounter from ray.data._internal.execution.interfaces.execution_options import ( ExecutionOptions, ExecutionResources, ) from ray.data._internal.execution.interfaces.physical_operator import ( PhysicalOperator, ReportsExtraResourceUsage, ) from ray.data._internal.execution.operators.base_physical_operator import ( AllToAllOperator, ) from ray.data._internal.execution.operators.hash_shuffle import ( HashShufflingOperatorBase, ) from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import ( # noqa: E501 ShuffleMapOp, ) from ray.data._internal.execution.operators.zip_operator import ZipOperator from ray.data._internal.execution.util import memory_string from ray.data.context import DataContext from ray.util.debug import log_once if TYPE_CHECKING: from ray.data._internal.execution.streaming_executor_state import OpState, Topology logger = logging.getLogger(__name__) LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE: Optional[bool] = env_bool( "RAY_DATA_DEBUG_RESOURCE_MANAGER", None ) # Only warn that the cluster can't run any task once the operator has been starved of # its minimum resources for this long. This avoids spurious warnings while the cluster # is still scaling up or waiting for a response from the autoscaling coordinator. # # I arbitrarily chose the default delay. STARVATION_WARNING_DELAY_S = env_float("RAY_DATA_STARVATION_WARNING_DELAY_S", 60) # Following list is a list of *blocking* materializing operators, that prevent # operators downstream from them from starting execution until these operators # finish executing. _BLOCKING_MATERIALIZING_OPERATORS = ( HashShufflingOperatorBase, AllToAllOperator, ShuffleMapOp, # TODO remove after zip made fully streaming ZipOperator, ) def terminal_operator_from_topology(topology: "Topology") -> PhysicalOperator: """Return the executor sink: the unique op with no in-DAG downstream consumers. ``build_streaming_topology`` is rooted at the same node passed to ``StreamingExecutor``; that root is the only operator whose ``output_dependencies`` is empty. """ if not topology: raise ValueError("topology must be non-empty") sinks = [op for op in topology if not op.output_dependencies] if len(sinks) == 1: return sinks[0] if not sinks: raise ValueError( "No terminal operator found in topology (expected exactly one " "operator with empty output_dependencies)" ) raise ValueError( "Expected exactly one terminal operator in topology, found " f"{len(sinks)}: {sinks!r}" ) class ResourceManager: """A class that manages the resource usage of a streaming executor.""" # The interval in seconds at which the global resource limits are refreshed. GLOBAL_LIMITS_UPDATE_INTERVAL_S = 1 # The fraction of the object store capacity that will be used as the default object # store memory limit for the streaming executor, # when `OpResourceAllocator` is enabled. DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION = env_float( "RAY_DATA_OBJECT_STORE_MEMORY_LIMIT_FRACTION", 0.5 ) # The fraction of the object store capacity that will be used as the default object # store memory limit for the streaming executor, # when `OpResourceAllocator` is not enabled. DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION = 0.25 def __init__( self, topology: "Topology", options: ExecutionOptions, get_total_resources: Callable[[], ExecutionResources], data_context: DataContext, block_ref_counter: BlockRefCounter, ): self._topology = topology self._options = options self._get_total_resources = get_total_resources self._global_limits = ExecutionResources.zero() self._global_limits_last_update_time = 0 self._global_usage = ExecutionResources.zero() self._global_running_usage = ExecutionResources.zero() self._global_pending_usage = ExecutionResources.zero() self._op_usages: Dict[PhysicalOperator, ExecutionResources] = {} self._op_running_usages: Dict[PhysicalOperator, ExecutionResources] = {} self._op_pending_usages: Dict[PhysicalOperator, ExecutionResources] = {} # Object store memory usage of pending task outputs (blocks being # generated by running tasks but not yet yielded). self._mem_op_internal: Dict[PhysicalOperator, int] = defaultdict(int) # Object store memory usage of the operator's outputs, including: # internal output queue, external output buffer in OpState, and # downstream operators' input buffers (inqueue + pending task inputs). self._mem_op_outputs: Dict[PhysicalOperator, int] = defaultdict(int) # Bytes buffered by external consumers (iterators) consuming Batches # (including the prefetched blocks). For example, # - ds.iter_batches -> one iterator # - streaming_split -> multiple iterators self._external_consumer_bytes: int = 0 self._has_external_consumer: bool = False # Executor sink (DAG root: unique op with no output_dependencies). # Iterator/streaming_split prefetch bytes are charged on this # operator's output usage. self._output_operator = terminal_operator_from_topology(topology) self._block_ref_counter = block_ref_counter self._op_resource_allocator: Optional[ "OpResourceAllocator" ] = create_resource_allocator(self, data_context) self._object_store_memory_limit_fraction = ( data_context.override_object_store_memory_limit_fraction if data_context.override_object_store_memory_limit_fraction is not None else ( self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION if self.op_resource_allocator_enabled() else self.DEFAULT_OBJECT_STORE_MEMORY_LIMIT_FRACTION_NO_RESERVATION ) ) @property def has_external_consumer(self) -> bool: """Return whether there is any external consumer.""" return self._has_external_consumer def set_external_consumer_bytes(self, num_bytes: int) -> None: """Set the bytes buffered by external consumers.""" assert ( num_bytes >= 0 ), f"external consumer bytes must be non-negative, got {num_bytes}" self._external_consumer_bytes = num_bytes self._has_external_consumer = True def get_external_consumer_bytes(self) -> int: """Get the bytes buffered by external consumers.""" return self._external_consumer_bytes def _estimate_object_store_memory_usage( self, op: "PhysicalOperator", state: "OpState" ) -> int: # Don't count input refs towards dynamic memory usage, as they have been # pre-created already outside this execution. if isinstance(op, InputDataBuffer): if op is self._output_operator: self._mem_op_internal[op] = 0 self._mem_op_outputs[op] = self._external_consumer_bytes return self._external_consumer_bytes return 0 usage = op.estimate_object_store_usage(state) self._mem_op_internal[op] = usage.internal self._mem_op_outputs[op] = usage.outputs # Attribute iterator / streaming_split prefetch to the executor sink only. if op is self._output_operator: self._mem_op_outputs[op] += self._external_consumer_bytes return self._mem_op_outputs[op] + self._mem_op_internal[op] def update_usages(self): """Recalculate resource usages.""" # TODO(hchen): This method will be called frequently during the execution loop. # And some computations are redundant. We should either remove redundant # computations or remove this method entirely and compute usages on demand. self._op_usages.clear() self._op_running_usages.clear() self._op_pending_usages.clear() # Iterate from last to first operator. for op, state in reversed(self._topology.items()): # Update `self._op_usages`, `self._op_running_usages`, # and `self._op_pending_usages`. op_usage = op.current_logical_usage() op_running_usage = op.running_logical_usage() op_pending_usage = op.pending_logical_usage() assert not op_usage.object_store_memory assert not op_running_usage.object_store_memory assert not op_pending_usage.object_store_memory used_object_store = self._estimate_object_store_memory_usage(op, state) op_usage = op_usage.copy(object_store_memory=used_object_store) op_running_usage = op_running_usage.copy( object_store_memory=used_object_store ) if isinstance(op, ReportsExtraResourceUsage): op_usage.add(op.extra_resource_usage()) self._op_usages[op] = op_usage self._op_running_usages[op] = op_running_usage self._op_pending_usages[op] = op_pending_usage # Update operator's object store usage, which is used by # DatasetStats and updated on the Ray Data dashboard. op._metrics.obj_store_mem_used = op_usage.object_store_memory # Roll the per-op usages up into the global totals in a single pass # each (one allocation per total instead of one per operator). self._global_usage = ExecutionResources.combine_sum(self._op_usages.values()) self._global_running_usage = ExecutionResources.combine_sum( self._op_running_usages.values() ) self._global_pending_usage = ExecutionResources.combine_sum( self._op_pending_usages.values() ) if self._op_resource_allocator is not None: self._update_allocated_budgets() def _update_allocated_budgets(self): completed_ops_usage = self._get_completed_ops_usage() available_limits = ( self.get_global_limits() .subtract(completed_ops_usage) .max(ExecutionResources.zero()) ) self._op_resource_allocator.update_budgets(limits=available_limits) def get_global_usage(self) -> ExecutionResources: """Return the global resource usage at the current time.""" assert ( self._global_usage.is_non_negative() ), f"Global usage should be non-negative, got {self._global_usage}" return self._global_usage def get_global_running_usage(self) -> ExecutionResources: """Return the global running resource usage at the current time.""" return self._global_running_usage def get_global_pending_usage(self) -> ExecutionResources: """Return the global pending resource usage at the current time.""" return self._global_pending_usage def get_global_limits(self) -> ExecutionResources: """Return the global resource limits at the current time. This method autodetects any unspecified execution resource limits based on the current cluster size, refreshing these values periodically to support cluster autoscaling. """ if ( time.time() - self._global_limits_last_update_time < self.GLOBAL_LIMITS_UPDATE_INTERVAL_S ): return self._global_limits self._global_limits_last_update_time = time.time() default_limits = self._options.resource_limits exclude = self._options.exclude_resources total_resources = self._get_total_resources() default_mem_fraction = self._object_store_memory_limit_fraction total_resources = total_resources.copy( object_store_memory=total_resources.object_store_memory * default_mem_fraction ) # Clamp to non-negative because exclude_resources (e.g., training worker # CPUs) can exceed the total resources reported by the cluster autoscaler, # such as when Ray Train reserves more CPUs than are visible to Ray Data. self._global_limits = ( default_limits.min(total_resources) .subtract(exclude) .max(ExecutionResources.zero()) ) return self._global_limits def get_op_usage( self, op: PhysicalOperator, include_ineligible_downstream: bool = False ) -> ExecutionResources: """Return the resource usage of the given operator at the current time.""" own_usage = self._op_usages[op] if not include_ineligible_downstream: return own_usage return own_usage.add(self._get_downstream_ineligible_ops_usage(op)) def _get_downstream_ineligible_ops_usage( self, op: PhysicalOperator ) -> ExecutionResources: return ExecutionResources.combine_sum( self.get_op_usage(op) for op in self._get_downstream_ineligible_ops(op) ) def get_mem_op_internal(self, op: PhysicalOperator) -> int: """Return the memory usage of pending task outputs for the given operator.""" return self._mem_op_internal[op] def get_mem_op_outputs( self, op: PhysicalOperator, include_ineligible_downstream: bool = False, ) -> int: """Return the memory usage of the outputs of the given operator.""" # Outputs usage of the current operator. op_outputs_usage = self._mem_op_outputs[op] if not include_ineligible_downstream: return op_outputs_usage # Also account the downstream ineligible operators' memory usage. return ( op_outputs_usage + self._get_downstream_ineligible_ops_usage(op).object_store_memory ) def get_op_usage_str(self, op: PhysicalOperator, *, verbose: bool) -> str: """Return a human-readable string representation of the resource usage of the given operator.""" # Handle case where operator is not in _op_running_usages dict if op not in self._op_running_usages: usage_str = "n/a" else: usage_str = f"{self._op_running_usages[op].cpu:.1f} CPU" if self._op_running_usages[op].memory: usage_str += f", {self._op_running_usages[op].memory_str()} memory" if self._op_running_usages[op].gpu: usage_str += f", {self._op_running_usages[op].gpu:.1f} GPU" usage_str += f", {self._op_running_usages[op].object_store_memory_str()} object store" # NOTE: Config can override requested verbosity level if LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE is not None: verbose = LOG_DEBUG_TELEMETRY_FOR_RESOURCE_MANAGER_OVERRIDE if verbose: usage_str += ( f" (in={memory_string(self.get_mem_op_internal(op))}," f"out={memory_string(self.get_mem_op_outputs(op))}" ) # External-consumer bytes (iterator / streaming_split prefetch) are # only attached to the output operator. Surface them in its line so # users can see how much of `out` is held by the downstream iterator # vs. the operator's own output queues. if op is self._output_operator and self._has_external_consumer: usage_str += ( f",external_consumer=" f"{memory_string(self._external_consumer_bytes)}" ) usage_str += ")" if self._op_resource_allocator is not None: allocation = self._op_resource_allocator.get_allocation(op) if allocation: usage_str += f", alloc=(cpu={allocation.cpu:.1f}" usage_str += f",mem={allocation.memory_str()}" usage_str += f",gpu={allocation.gpu:.1f}" usage_str += f",obj_store={allocation.object_store_memory_str()})" budget = self._op_resource_allocator.get_budget(op) if budget: usage_str += f", budget=(cpu={budget.cpu:.1f}" usage_str += f",mem={budget.memory_str()}" usage_str += f",gpu={budget.gpu:.1f}" usage_str += f",obj_store={budget.object_store_memory_str()}" # Remaining memory budget for producing new task outputs. if isinstance( self._op_resource_allocator, ReservationOpResourceAllocator ): reserved_for_output = memory_string( self._op_resource_allocator._output_budgets.get(op, 0) ) usage_str += f",out={reserved_for_output})" return usage_str def op_resource_allocator_enabled(self) -> bool: """Return whether OpResourceAllocator is enabled.""" return self._op_resource_allocator is not None @property def op_resource_allocator(self) -> "OpResourceAllocator": """Return the OpResourceAllocator.""" assert self._op_resource_allocator is not None return self._op_resource_allocator def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]: """Return the budget for the given operator, or None if the operator has unlimited budget.""" if self._op_resource_allocator is None: return None return self._op_resource_allocator.get_budget(op) def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]: """Return the allocation of the given operator, or None if the operator doesn't have a designated allocation.""" if self._op_resource_allocator is None: return None return self._op_resource_allocator.get_allocation(op) def is_op_eligible(self, op: PhysicalOperator) -> bool: """Whether the op is eligible for memory reservation.""" return ( not op.throttling_disabled() # As long as the op has finished execution, even if there are still # non-taken outputs, we don't need to allocate resources for it. and not op.has_execution_finished() ) def _get_downstream_ineligible_ops( self, op: PhysicalOperator ) -> Iterable[PhysicalOperator]: """Get the downstream ineligible operators of the given operator. E.g., - "cur_map->downstream_map" will return an empty list. - "cur_map->limit1->limit2->downstream_map" will return [limit1, limit2]. """ for next_op in op.output_dependencies: if not self.is_op_eligible(next_op): yield next_op yield from self._get_downstream_ineligible_ops(next_op) def get_downstream_eligible_ops( self, op: PhysicalOperator ) -> Iterable[PhysicalOperator]: """Get the downstream eligible operators of the given operator, ignoring intermediate ineligible operators. E.g., - "cur_map->downstream_map" will return [downstream_map]. - "cur_map->limit1->limit2->downstream_map" will return [downstream_map]. """ for next_op in op.output_dependencies: if self.is_op_eligible(next_op): yield next_op else: yield from self.get_downstream_eligible_ops(next_op) def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]: if self._op_resource_allocator is not None: return self._op_resource_allocator.max_task_output_bytes_to_read(op) return None def _get_completed_ops_usage(self) -> ExecutionResources: """ Resource reservation is based on the number of eligible operators. However, there might be completed operators that still have blocks in their output queue, which we need to exclude them from the reservation. And we also need to exclude the downstream ineligible operators. E.g., for the following pipeline: ``` map1 (completed, but still has blocks in its output queue) -> limit1 (ineligible, not completed) -> map2 (not completed) -> limit2 -> map3 ``` The reservation is based on the number of eligible operators (map2 and map3), but we need to exclude map1 and limit1 from the reservation. """ last_completed_ops = [] ops_to_exclude = [] # Traverse operator tree collecting all operators that have already finished for op in self._topology: if not op.has_execution_finished(): for dep in op.input_dependencies: if dep.has_execution_finished(): last_completed_ops.append(dep) # In addition to completed operators, # filter out downstream ineligible operators since they are omitted from reservation calculations. for op in last_completed_ops: ops_to_exclude.extend(list(self._get_downstream_ineligible_ops(op))) ops_to_exclude.append(op) completed_ops = list(set(ops_to_exclude)) completed_ops_usage = ExecutionResources.combine_sum( self.get_op_usage(op) for op in completed_ops ) return completed_ops_usage def _is_blocking_materializing_op(self, op: PhysicalOperator) -> bool: """This method checks whether either 1. Operator itself or 2. One of operator's immediate *ineligible* downstream dependencies Are blocking, materializing operators. NOTE: That downstream ineligible operators are considered an "extension" of their first preceding eligible operator from resource allocation standpoint. """ # Check if Op itself is a blocking, materializing operator if isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS): return True # Check if any of its direct *ineligible* downstream dependencies are # blocking, materializing operators. # # NOTE: We only check ineligible downstream deps, since eligible downstream # deps will have their own allocation that is adjusted appropriately return any( isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS) for op in self._get_downstream_ineligible_ops(op) ) def _get_first_pending_materializing_op(topology: "Topology") -> int: for idx, op in enumerate(topology): if isinstance(op, _BLOCKING_MATERIALIZING_OPERATORS) and not op.has_completed(): return idx return -1 class OpResourceAllocator(ABC): """An interface for dynamic operator resource allocation. This interface allows dynamically allocating available resources to each operator, limiting how many tasks each operator can submit, and how much data each operator can read from its running tasks. """ def __init__(self, resource_manager: "ResourceManager"): self._resource_manager = resource_manager self._topology = resource_manager._topology @abstractmethod def update_budgets( self, *, limits: ExecutionResources, ): """Callback to update resource usages.""" ... @abstractmethod def can_submit_new_task(self, op: PhysicalOperator) -> bool: """Return whether the given operator can submit a new task.""" ... @abstractmethod def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]: """Return the maximum bytes of pending task outputs can be read for the given operator. None means no limit.""" ... @abstractmethod def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]: """Returns the budget for the given operator or `None` if the operator has unlimited budget. Operator's budget is defined as: Budget = Allocation - Usage """ ... @abstractmethod def get_output_budget(self, op: PhysicalOperator) -> Optional[int]: """Returns the budget for operator's outputs (in object store bytes) or `None` if there's no limit. """ ... @abstractmethod def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]: """Returns allocation for the given operator or `None` if operator's allocation is unlimited.""" ... def _get_eligible_ops(self) -> List[PhysicalOperator]: """Returns a list of operators eligible for allocation. Only operators upstream of the first non-completed *materializing* Operator, (like shuffle, etc) are able to receive inputs and start execution. Therefore only these operators are eligible for resource allocation. """ first_pending_materializing_op_idx = _get_first_pending_materializing_op( self._topology ) return [ op for idx, op in enumerate(self._topology) if self._is_op_eligible(op) and ( first_pending_materializing_op_idx == -1 or idx <= first_pending_materializing_op_idx ) ] @staticmethod def _is_op_eligible(op: PhysicalOperator) -> bool: """Whether the op is eligible for memory reservation.""" return ( not op.throttling_disabled() # As long as the op has finished execution, even if there are still # non-taken outputs, we don't need to allocate resources for it. and not op.has_execution_finished() ) class ReservationOpResourceAllocator(OpResourceAllocator): """An OpResourceAllocator implementation that reserves resources for each operator. This class reserves memory and CPU resources for eligible operators, and considers runtime resource usages to limit the resources that each operator can use. It works in the following way: 1. An operator is eligible for resource reservation, if it has enabled throttling and hasn't completed. Ineligible operators are not throttled, but their usage will be accounted for their upstream eligible operators. E.g., for such a dataset "map1->limit->map2->streaming_split", we'll treat "map1->limit" as a group and "map2->streaming_split" as another group. 2. For each eligible operator, we reserve `reservation_ratio * global_resources / num_eligible_ops` resources, half of which is reserved only for the operator outputs, excluding pending task outputs. 3. Non-reserved resources are shared among all operators. 4. In each scheduling iteration, each eligible operator will get "remaining of their own reserved resources" + "remaining of shared resources / num_eligible_ops" resources. The `reservation_ratio` is set to 50% by default. Users can tune this value to adjust how aggressive or conservative the resource allocation is. A higher value will make the resource allocation more even, but may lead to underutilization and worse performance. And vice versa. """ def __init__(self, resource_manager: ResourceManager, reservation_ratio: float): super().__init__(resource_manager) self._reservation_ratio = reservation_ratio assert 0.0 <= self._reservation_ratio <= 1.0 # Per-op reserved resources, excluding `_reserved_for_op_outputs`. self._op_reserved: Dict[PhysicalOperator, ExecutionResources] = {} # Memory reserved exclusively for the outputs of each operator. # "Op outputs" refer to blocks that have been taken out of an operator, # i.e., `RessourceManager._mem_op_outputs`. # # Note, if we don't reserve memory for op outputs, all the budget may be used by # the pending task outputs, and/or op's internal output buffers (the latter can # happen when `preserve_order=True`). # Then we'll have no budget to pull blocks from the op. self._reserved_for_op_outputs: Dict[PhysicalOperator, float] = {} # Total shared resources. self._total_shared = ExecutionResources.zero() # Resource budgets for each operator, excluding `_reserved_for_op_outputs`. self._op_budgets: Dict[PhysicalOperator, ExecutionResources] = {} # Remaining memory budget for generating new task outputs, per operator. self._output_budgets: Dict[PhysicalOperator, float] = {} # Whether each operator has reserved the minimum resources to run # at least one task. # This is used to avoid edge cases where the entire resource limits are not # enough to run one task of each op. # See `test_no_deadlock_on_small_cluster_resources` as an example. self._reserved_min_resources: Dict[PhysicalOperator, bool] = {} # `time.monotonic()` timestamp at which each operator most recently became # starved of its minimum resources, or None if it currently has them. self._op_starved_since: Dict[PhysicalOperator, Optional[float]] = {} def _update_reservation(self, limits: ExecutionResources): eligible_ops = self._get_eligible_ops() self._op_reserved.clear() self._reserved_for_op_outputs.clear() self._reserved_min_resources.clear() if len(eligible_ops) == 0: return remaining = limits.copy() # Reserve `reservation_ratio * global_limits / num_ops` resources for each # operator. default_reserved = limits.scale(self._reservation_ratio / (len(eligible_ops))) for index, op in enumerate(eligible_ops): # Reserve at least half of the default reserved resources for the outputs. # This makes sure that we will have enough budget to pull blocks from the # op. reserved_for_outputs = ExecutionResources( 0, 0, max(default_reserved.object_store_memory / 2, 1) ) reserved_for_tasks = default_reserved.subtract(reserved_for_outputs) min_resource_usage, max_resource_usage = op.min_max_resource_requirements() if min_resource_usage is not None: reserved_for_tasks = reserved_for_tasks.max(min_resource_usage) if max_resource_usage is not None: reserved_for_tasks = reserved_for_tasks.min(max_resource_usage) # Check if the remaining resources are enough for both reserved_for_tasks # and reserved_for_outputs. Note, we only consider CPU and GPU, but not # object_store_memory, because object_store_memory can be oversubscribed, # but CPU/GPU cannot. if reserved_for_tasks.add(reserved_for_outputs).satisfies_limit( remaining, ignore_object_store_memory=True ): self._reserved_min_resources[op] = True self._op_starved_since[op] = None else: self._reserved_min_resources[op] = False if self._op_starved_since.get(op) is None: self._op_starved_since[op] = time.monotonic() # If the remaining resources are not enough to reserve the minimum # resources for this operator, we'll only reserve the minimum object # store memory, but not the CPU and GPU resources. # Because Ray Core doesn't allow CPU/GPU resources to be oversubscribed. # NOTE: we prioritize upstream operators for minimum resource reservation. # ops. It's fine that downstream ops don't get the minimum reservation, # because they can wait for upstream ops to finish and release resources. reserved_for_tasks = ExecutionResources( 0, 0, min_resource_usage.object_store_memory ) # Log a warning if even the first operator cannot reserve the minimum # resources. if index == 0: self._warn_if_op_starved_too_long(op) self._op_reserved[op] = reserved_for_tasks self._reserved_for_op_outputs[op] = reserved_for_outputs.object_store_memory op_total_reserved = reserved_for_tasks.add(reserved_for_outputs) remaining = remaining.subtract(op_total_reserved) remaining = remaining.max(ExecutionResources.zero()) self._total_shared = remaining def _warn_if_op_starved_too_long(self, op: PhysicalOperator) -> None: # The operator isn't starved. Return early. if self._op_starved_since.get(op) is None: return op_starved_duration = time.monotonic() - self._op_starved_since[op] if ( op_starved_duration >= STARVATION_WARNING_DELAY_S # Add `id(self)` to the log_once key so that it will be logged once per # execution. and log_once(f"starvation_warning_{id(self)}") ): logger.warning( f"Cluster resources are not enough to run any task from {op}." " The job may hang forever unless the cluster scales up." ) def can_submit_new_task(self, op: PhysicalOperator) -> bool: """Return whether the given operator can submit a new task based on budget.""" budget = self.get_budget(op) if budget is None: return True return ( op.incremental_resource_usage().satisfies_limit(budget) and # Avoid scheduling if there's no more Object Store budget (for # task outputs) budget.object_store_memory >= (op.metrics.obj_store_mem_max_pending_output_per_task or 0) ) def get_budget(self, op: PhysicalOperator) -> Optional[ExecutionResources]: return self._op_budgets.get(op) def get_output_budget(self, op: PhysicalOperator) -> Optional[int]: return self._output_budgets.get(op) def get_allocation(self, op: PhysicalOperator) -> Optional[ExecutionResources]: budget = self.get_budget(op) if budget is None: return None return budget.add(self._resource_manager.get_op_usage(op)) def _get_total_reserved(self, op: PhysicalOperator) -> ExecutionResources: """Get total reserved resources for an operator, including outputs reservation.""" op_reserved = self._op_reserved[op] reserved_for_outputs = self._reserved_for_op_outputs[op] return op_reserved.copy( object_store_memory=op_reserved.object_store_memory + reserved_for_outputs ) def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]: if op not in self._op_budgets: return None res = self._op_budgets[op].object_store_memory # Add the remaining of `_reserved_for_op_outputs`. op_outputs_usage = self._resource_manager.get_mem_op_outputs( op, include_ineligible_downstream=True ) res += max(self._reserved_for_op_outputs[op] - op_outputs_usage, 0) if math.isinf(res): self._output_budgets[op] = res return None res = int(res) assert res >= 0 self._output_budgets[op] = res return res def update_budgets( self, *, limits: ExecutionResources, ): # Remaining resources to be distributed across operators remaining_shared = self._update_reservation(limits) self._op_budgets.clear() eligible_ops = self._get_eligible_ops() if len(eligible_ops) == 0: return # Remaining of shared resources. remaining_shared = self._total_shared for op in eligible_ops: # Calculate the memory usage of the operator. op_mem_usage = 0 # Add the memory usage of the operator itself, # excluding `_reserved_for_op_outputs`. op_mem_usage += self._resource_manager.get_mem_op_internal(op) # Add the portion of op outputs usage that has # exceeded `_reserved_for_op_outputs`. op_outputs_usage = self._resource_manager.get_mem_op_outputs( op, include_ineligible_downstream=True ) op_mem_usage += max(op_outputs_usage - self._reserved_for_op_outputs[op], 0) op_usage = self._resource_manager.get_op_usage(op).copy( object_store_memory=op_mem_usage ) op_reserved = self._op_reserved[op] # How much of the reserved resources are remaining. op_reserved_remaining = op_reserved.subtract(op_usage).max( ExecutionResources.zero() ) self._op_budgets[op] = op_reserved_remaining # How much of the reserved resources are exceeded. # If exceeded, we need to subtract from the remaining shared resources. op_reserved_exceeded = op_usage.subtract(op_reserved).max( ExecutionResources.zero() ) remaining_shared = remaining_shared.subtract(op_reserved_exceeded) remaining_shared = remaining_shared.max(ExecutionResources.zero()) # Allocate the remaining shared resources to each operator. for i, op in enumerate(reversed(eligible_ops)): # By default, divide the remaining shared resources equally. op_shared = remaining_shared.scale(1.0 / (len(eligible_ops) - i)) # But if the op's budget is less than `min_scheduling_resources`, # it will be useless. So we'll let the downstream operator # borrow some resources from the upstream operator, if remaining_shared # is still enough. to_borrow = ( op.min_scheduling_resources() .subtract(self._op_budgets[op].add(op_shared)) .max(ExecutionResources.zero()) ) if not to_borrow.is_zero() and op_shared.add(to_borrow).satisfies_limit( remaining_shared ): op_shared = op_shared.add(to_borrow) # Cap op_shared so that total allocation doesn't exceed max_resource_usage. # Total allocation = max(total_reserved, op_usage) + op_shared # This ensures excess resources stay in remaining_shared for other operators. _, max_resource_usage = op.min_max_resource_requirements() if max_resource_usage != ExecutionResources.inf(): total_reserved = self._get_total_reserved(op) op_usage = self._resource_manager.get_op_usage(op) current_allocation = total_reserved.max(op_usage) max_shared = max_resource_usage.subtract(current_allocation).max( ExecutionResources.zero() ) op_shared = op_shared.min(max_shared) remaining_shared = remaining_shared.subtract(op_shared) assert remaining_shared.is_non_negative(), ( remaining_shared, op, op_shared, to_borrow, ) self._op_budgets[op] = self._op_budgets[op].add(op_shared) # Give any remaining shared resources to the most downstream uncapped op. # This can happen when some ops have their shared allocation capped. if eligible_ops and not remaining_shared.is_zero(): for op in reversed(eligible_ops): _, max_resource_usage = op.min_max_resource_requirements() if max_resource_usage == ExecutionResources.inf(): self._op_budgets[op] = self._op_budgets[op].add(remaining_shared) break # A materializing operator like `AllToAllOperator` waits for all its input # operator's outputs before processing data. This often forces the input # operator to exceed its object store memory budget. To prevent deadlock, we # disable object store memory backpressure for the input operator. for op in eligible_ops: if self._resource_manager._is_blocking_materializing_op(op): self._op_budgets[op] = self._op_budgets[op].copy( object_store_memory=float("inf") )